From a7a6fe261d679dfbe8cf6c5e90c7f5981ff08abd Mon Sep 17 00:00:00 2001
From: epenet <6771947+epenet@users.noreply.github.com>
Date: Wed, 4 Sep 2024 07:08:50 +0000
Subject: [PATCH] Improve config flow type hints in screenlogic

---
 .../components/screenlogic/config_flow.py        | 16 +++++++++-------
 .../components/screenlogic/coordinator.py        |  5 ++++-
 2 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/homeassistant/components/screenlogic/config_flow.py b/homeassistant/components/screenlogic/config_flow.py
index 74a01fdeaa2ae1..4a46756cf2f6a1 100644
--- a/homeassistant/components/screenlogic/config_flow.py
+++ b/homeassistant/components/screenlogic/config_flow.py
@@ -32,9 +32,9 @@
 PENTAIR_OUI = "00-C0-33"
 
 
-async def async_discover_gateways_by_unique_id(hass):
+async def async_discover_gateways_by_unique_id() -> dict[str, dict[str, Any]]:
     """Discover gateways and return a dict of them by unique id."""
-    discovered_gateways = {}
+    discovered_gateways: dict[str, dict[str, Any]] = {}
     try:
         hosts = await discovery.async_discover()
         _LOGGER.debug("Discovered hosts: %s", hosts)
@@ -51,16 +51,16 @@ async def async_discover_gateways_by_unique_id(hass):
     return discovered_gateways
 
 
-def _extract_mac_from_name(name):
+def _extract_mac_from_name(name: str) -> str:
     return format_mac(f"{PENTAIR_OUI}-{name.split(':')[1].strip()}")
 
 
-def short_mac(mac):
+def short_mac(mac: str) -> str:
     """Short version of the mac as seen in the app."""
     return "-".join(mac.split(":")[3:]).upper()
 
 
-def name_for_mac(mac):
+def name_for_mac(mac: str) -> str:
     """Derive the gateway name from the mac."""
     return f"Pentair: {short_mac(mac)}"
 
@@ -83,9 +83,11 @@ def async_get_options_flow(
         """Get the options flow for ScreenLogic."""
         return ScreenLogicOptionsFlowHandler(config_entry)
 
-    async def async_step_user(self, user_input=None) -> ConfigFlowResult:
+    async def async_step_user(
+        self, user_input: dict[str, Any] | None = None
+    ) -> ConfigFlowResult:
         """Handle the start of the config flow."""
-        self.discovered_gateways = await async_discover_gateways_by_unique_id(self.hass)
+        self.discovered_gateways = await async_discover_gateways_by_unique_id()
         return await self.async_step_gateway_select()
 
     async def async_step_dhcp(
diff --git a/homeassistant/components/screenlogic/coordinator.py b/homeassistant/components/screenlogic/coordinator.py
index 281bac86e01025..a90c9cb2cf40d5 100644
--- a/homeassistant/components/screenlogic/coordinator.py
+++ b/homeassistant/components/screenlogic/coordinator.py
@@ -2,6 +2,7 @@
 
 from datetime import timedelta
 import logging
+from typing import TYPE_CHECKING
 
 from screenlogicpy import ScreenLogicGateway
 from screenlogicpy.const.common import (
@@ -33,11 +34,13 @@ async def async_get_connect_info(
     """Construct connect_info from configuration entry and returns it to caller."""
     mac = entry.unique_id
     # Attempt to rediscover gateway to follow IP changes
-    discovered_gateways = await async_discover_gateways_by_unique_id(hass)
+    discovered_gateways = await async_discover_gateways_by_unique_id()
     if mac in discovered_gateways:
         return discovered_gateways[mac]
 
     _LOGGER.debug("Gateway rediscovery failed for %s", entry.title)
+    if TYPE_CHECKING:
+        assert mac is not None
     # Static connection defined or fallback from discovery
     return {
         SL_GATEWAY_NAME: name_for_mac(mac),