Improve config flow type hints in yeelight (#125319)

homeassistant/components/yeelight/config_flow.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 import logging
-from typing import TYPE_CHECKING, Any
+from typing import Any
 from urllib.parse import urlparse
 
 import voluptuous as vol
@@ -23,6 +23,7 @@ from homeassistant.const import CONF_DEVICE, CONF_HOST, CONF_ID, CONF_MODEL, CON
 from homeassistant.core import callback
 from homeassistant.exceptions import HomeAssistantError
 import homeassistant.helpers.config_validation as cv
+from homeassistant.helpers.typing import VolDictType
 
 from .const import (
     CONF_DETECTED_MODEL,
@@ -52,6 +53,9 @@ class YeelightConfigFlow(ConfigFlow, domain=DOMAIN):
 
     VERSION = 1
 
+    _discovered_ip: str
+    _discovered_model: str
+
     @staticmethod
     @callback
     def async_get_options_flow(config_entry: ConfigEntry) -> OptionsFlowHandler:
@@ -61,8 +65,6 @@ class YeelightConfigFlow(ConfigFlow, domain=DOMAIN):
     def __init__(self) -> None:
         """Initialize the config flow."""
         self._discovered_devices: dict[str, Any] = {}
-        self._discovered_model = None
-        self._discovered_ip: str | None = None
 
     async def async_step_homekit(
         self, discovery_info: zeroconf.ZeroconfServiceInfo
@@ -96,7 +98,7 @@ class YeelightConfigFlow(ConfigFlow, domain=DOMAIN):
         await self.async_set_unique_id(discovery_info.ssdp_headers["id"])
         return await self._async_handle_discovery_with_unique_id()
 
-    async def _async_handle_discovery_with_unique_id(self):
+    async def _async_handle_discovery_with_unique_id(self) -> ConfigFlowResult:
         """Handle any discovery with a unique id."""
         for entry in self._async_current_entries(include_ignore=False):
             if entry.unique_id != self.unique_id and self.unique_id != entry.data.get(
@@ -117,7 +119,7 @@ class YeelightConfigFlow(ConfigFlow, domain=DOMAIN):
             return self.async_abort(reason="already_configured")
         return await self._async_handle_discovery()
 
-    async def _async_handle_discovery(self):
+    async def _async_handle_discovery(self) -> ConfigFlowResult:
         """Handle any discovery."""
         self.context[CONF_HOST] = self._discovered_ip
         for progress in self._async_in_progress():
@@ -140,7 +142,9 @@ class YeelightConfigFlow(ConfigFlow, domain=DOMAIN):
         )
         return await self.async_step_discovery_confirm()
 
-    async def async_step_discovery_confirm(self, user_input=None):
+    async def async_step_discovery_confirm(
+        self, user_input: dict[str, Any] | None = None
+    ) -> ConfigFlowResult:
         """Confirm discovery."""
         if user_input is not None or not onboarding.async_is_onboarded(self.hass):
             return self.async_create_entry(
@@ -179,8 +183,6 @@ class YeelightConfigFlow(ConfigFlow, domain=DOMAIN):
                 errors["base"] = "cannot_connect"
             else:
                 self._abort_if_unique_id_configured()
-                if TYPE_CHECKING:
-                    assert self.unique_id
                 return self.async_create_entry(
                     title=async_format_model_id(model, self.unique_id),
                     data={
@@ -199,7 +201,9 @@ class YeelightConfigFlow(ConfigFlow, domain=DOMAIN):
             errors=errors,
         )
 
-    async def async_step_pick_device(self, user_input=None):
+    async def async_step_pick_device(
+        self, user_input: dict[str, str] | None = None
+    ) -> ConfigFlowResult:
         """Handle the step to pick discovered device."""
         if user_input is not None:
             unique_id = user_input[CONF_DEVICE]
@@ -260,7 +264,9 @@ class YeelightConfigFlow(ConfigFlow, domain=DOMAIN):
         self._abort_if_unique_id_configured()
         return self.async_create_entry(title=import_data[CONF_NAME], data=import_data)
 
-    async def _async_try_connect(self, host, raise_on_progress=True):
+    async def _async_try_connect(
+        self, host: str, raise_on_progress: bool = True
+    ) -> str:
         """Set up with options."""
         self._async_abort_entries_match({CONF_HOST: host})
 
@@ -294,7 +300,9 @@ class OptionsFlowHandler(OptionsFlow):
         """Initialize the option flow."""
         self._config_entry = config_entry
 
-    async def async_step_init(self, user_input=None):
+    async def async_step_init(
+        self, user_input: dict[str, Any] | None = None
+    ) -> ConfigFlowResult:
         """Handle the initial step."""
         data = self._config_entry.data
         options = self._config_entry.options
@@ -306,7 +314,7 @@ class OptionsFlowHandler(OptionsFlow):
                 title="", data={CONF_MODEL: model, **options, **user_input}
             )
 
-        schema_dict = {}
+        schema_dict: VolDictType = {}
         known_models = get_known_models()
         if is_unknown_model := model not in known_models:
             known_models.insert(0, model)

homeassistant/components/yeelight/device.py

@@ -32,13 +32,13 @@ def async_format_model(model: str) -> str:
 
 
 @callback
-def async_format_id(id_: str) -> str:
+def async_format_id(id_: str | None) -> str:
     """Generate a more human readable id."""
     return hex(int(id_, 16)) if id_ else "None"
 
 
 @callback
-def async_format_model_id(model: str, id_: str) -> str:
+def async_format_model_id(model: str, id_: str | None) -> str:
     """Generate a more human readable name."""
     return f"{async_format_model(model)} {async_format_id(id_)}"