Second part of Strict types for Fritz (#52086)

Co-authored-by: Shay Levy <levyshay1@gmail.com>
Co-authored-by: Martin Hjelmare <marhje52@gmail.com>

homeassistant/components/fritz/binary_sensor.py

@@ -24,7 +24,7 @@ async def async_setup_entry(
     _LOGGER.debug("Setting up FRITZ!Box binary sensors")
     fritzbox_tools: FritzBoxTools = hass.data[DOMAIN][entry.entry_id]
 
-    if "WANIPConn1" in fritzbox_tools.connection.services:
+    if fritzbox_tools.connection and "WANIPConn1" in fritzbox_tools.connection.services:
         # Only routers are supported at the moment
         async_add_entities(
             [FritzBoxConnectivitySensor(fritzbox_tools, entry.title)], True
@@ -74,14 +74,19 @@ class FritzBoxConnectivitySensor(FritzBoxBaseEntity, BinarySensorEntity):
         _LOGGER.debug("Updating FRITZ!Box binary sensors")
         self._is_on = True
         try:
-            if "WANCommonInterfaceConfig1" in self._fritzbox_tools.connection.services:
+            if (
+                self._fritzbox_tools.connection
+                and "WANCommonInterfaceConfig1"
+                in self._fritzbox_tools.connection.services
+            ):
                 link_props = self._fritzbox_tools.connection.call_action(
                     "WANCommonInterfaceConfig1", "GetCommonLinkProperties"
                 )
                 is_up = link_props["NewPhysicalLinkStatus"]
                 self._is_on = is_up == "Up"
             else:
-                self._is_on = self._fritzbox_tools.fritz_status.is_connected
+                if self._fritzbox_tools.fritz_status:
+                    self._is_on = self._fritzbox_tools.fritz_status.is_connected
 
             self._is_available = True
 

homeassistant/components/fritz/common.py

@@ -1,12 +1,12 @@
 """Support for AVM FRITZ!Box classes."""
 from __future__ import annotations
 
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from datetime import datetime, timedelta
 import logging
-from typing import Any
+from types import MappingProxyType
+from typing import Any, TypedDict
 
-# pylint: disable=import-error
 from fritzconnection import FritzConnection
 from fritzconnection.core.exceptions import (
     FritzActionError,
@@ -20,10 +20,11 @@ from homeassistant.components.device_tracker.const import (
     CONF_CONSIDER_HOME,
     DEFAULT_CONSIDER_HOME,
 )
-from homeassistant.core import callback
+from homeassistant.core import CALLBACK_TYPE, HomeAssistant, callback
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC
 from homeassistant.helpers.dispatcher import dispatcher_send
+from homeassistant.helpers.entity import DeviceInfo
 from homeassistant.helpers.event import async_track_time_interval
 from homeassistant.util import dt as dt_util
 
@@ -40,6 +41,14 @@ from .const import (
 _LOGGER = logging.getLogger(__name__)
 
 
+class ClassSetupMissing(Exception):
+    """Raised when a Class func is called before setup."""
+
+    def __init__(self):
+        """Init custom exception."""
+        super().__init__("Function called before Class setup")
+
+
 @dataclass
 class Device:
     """FRITZ!Box device class."""
@@ -49,39 +58,48 @@ class Device:
     name: str
 
 
+class HostInfo(TypedDict):
+    """FRITZ!Box host info class."""
+
+    mac: str
+    name: str
+    ip: str
+    status: bool
+
+
 class FritzBoxTools:
     """FrtizBoxTools class."""
 
     def __init__(
         self,
-        hass,
-        password,
-        username=DEFAULT_USERNAME,
-        host=DEFAULT_HOST,
-        port=DEFAULT_PORT,
-    ):
+        hass: HomeAssistant,
+        password: str,
+        username: str = DEFAULT_USERNAME,
+        host: str = DEFAULT_HOST,
+        port: int = DEFAULT_PORT,
+    ) -> None:
         """Initialize FritzboxTools class."""
-        self._cancel_scan = None
+        self._cancel_scan: CALLBACK_TYPE | None = None
         self._devices: dict[str, Any] = {}
-        self._options = None
-        self._unique_id = None
-        self.connection = None
-        self.fritz_hosts = None
-        self.fritz_status = None
+        self._options: MappingProxyType[str, Any] | None = None
+        self._unique_id: str | None = None
+        self.connection: FritzConnection = None
+        self.fritz_hosts: FritzHosts = None
+        self.fritz_status: FritzStatus = None
         self.hass = hass
         self.host = host
         self.password = password
         self.port = port
         self.username = username
-        self.mac = None
-        self.model = None
-        self.sw_version = None
+        self._mac: str | None = None
+        self._model: str | None = None
+        self._sw_version: str | None = None
 
-    async def async_setup(self):
+    async def async_setup(self) -> None:
         """Wrap up FritzboxTools class setup."""
-        return await self.hass.async_add_executor_job(self.setup)
+        await self.hass.async_add_executor_job(self.setup)
 
-    def setup(self):
+    def setup(self) -> None:
         """Set up FritzboxTools class."""
         self.connection = FritzConnection(
             address=self.host,
@@ -93,14 +111,13 @@ class FritzBoxTools:
 
         self.fritz_status = FritzStatus(fc=self.connection)
         info = self.connection.call_action("DeviceInfo:1", "GetInfo")
-        if self._unique_id is None:
+        if not self._unique_id:
             self._unique_id = info["NewSerialNumber"]
 
-        self.model = info.get("NewModelName")
-        self.sw_version = info.get("NewSoftwareVersion")
-        self.mac = self.unique_id
+        self._model = info.get("NewModelName")
+        self._sw_version = info.get("NewSoftwareVersion")
 
-    async def async_start(self, options):
+    async def async_start(self, options: MappingProxyType[str, Any]) -> None:
         """Start FritzHosts connection."""
         self.fritz_hosts = FritzHosts(fc=self.connection)
         self._options = options
@@ -111,7 +128,7 @@ class FritzBoxTools:
         )
 
     @callback
-    def async_unload(self):
+    def async_unload(self) -> None:
         """Unload FritzboxTools class."""
         _LOGGER.debug("Unloading FRITZ!Box router integration")
         if self._cancel_scan is not None:
@@ -119,8 +136,31 @@ class FritzBoxTools:
             self._cancel_scan = None
 
     @property
-    def unique_id(self):
+    def unique_id(self) -> str:
         """Return unique id."""
+        if not self._unique_id:
+            raise ClassSetupMissing()
+        return self._unique_id
+
+    @property
+    def model(self) -> str:
+        """Return device model."""
+        if not self._model:
+            raise ClassSetupMissing()
+        return self._model
+
+    @property
+    def sw_version(self) -> str:
+        """Return SW version."""
+        if not self._sw_version:
+            raise ClassSetupMissing()
+        return self._sw_version
+
+    @property
+    def mac(self) -> str:
+        """Return device Mac address."""
+        if not self._unique_id:
+            raise ClassSetupMissing()
         return self._unique_id
 
     @property
@@ -138,7 +178,7 @@ class FritzBoxTools:
         """Event specific per FRITZ!Box entry to signal updates in devices."""
         return f"{DOMAIN}-device-update-{self._unique_id}"
 
-    def _update_info(self):
+    def _update_info(self) -> list[HostInfo]:
         """Retrieve latest information from the FRITZ!Box."""
         return self.fritz_hosts.get_hosts_info()
 
@@ -146,9 +186,12 @@ class FritzBoxTools:
         """Scan for new devices and return a list of found device ids."""
         _LOGGER.debug("Checking devices for FRITZ!Box router %s", self.host)
 
-        consider_home = self._options.get(
-            CONF_CONSIDER_HOME, DEFAULT_CONSIDER_HOME.total_seconds()
-        )
+        if self._options:
+            consider_home = self._options.get(
+                CONF_CONSIDER_HOME, DEFAULT_CONSIDER_HOME.total_seconds()
+            )
+        else:
+            consider_home = DEFAULT_CONSIDER_HOME
 
         new_device = False
         for known_host in self._update_info():
@@ -165,7 +208,7 @@ class FritzBoxTools:
             if dev_mac in self._devices:
                 self._devices[dev_mac].update(dev_info, dev_home, consider_home)
             else:
-                device = FritzDevice(dev_mac)
+                device = FritzDevice(dev_mac, dev_name)
                 device.update(dev_info, dev_home, consider_home)
                 self._devices[dev_mac] = device
                 new_device = True
@@ -177,6 +220,10 @@ class FritzBoxTools:
     async def service_fritzbox(self, service: str) -> None:
         """Define FRITZ!Box services."""
         _LOGGER.debug("FRITZ!Box router: %s", service)
+
+        if not self.connection:
+            raise HomeAssistantError("Unable to establish a connection")
+
         try:
             if service == SERVICE_REBOOT:
                 await self.hass.async_add_executor_job(
@@ -194,26 +241,25 @@ class FritzBoxTools:
             raise HomeAssistantError("Service not supported") from ex
 
 
+@dataclass
 class FritzData:
     """Storage class for platform global data."""
 
-    def __init__(self) -> None:
-        """Initialize the data."""
-        self.tracked: dict = {}
+    tracked: dict = field(default_factory=dict)
 
 
 class FritzDevice:
     """FritzScanner device."""
 
-    def __init__(self, mac, name=None):
+    def __init__(self, mac: str, name: str) -> None:
         """Initialize device info."""
         self._mac = mac
         self._name = name
-        self._ip_address = None
-        self._last_activity = None
+        self._ip_address: str | None = None
+        self._last_activity: datetime | None = None
         self._connected = False
 
-    def update(self, dev_info, dev_home, consider_home):
+    def update(self, dev_info: Device, dev_home: bool, consider_home: float) -> None:
         """Update device info."""
         utc_point_in_time = dt_util.utcnow()
 
@@ -235,27 +281,27 @@ class FritzDevice:
         self._ip_address = dev_info.ip_address if self._connected else None
 
     @property
-    def is_connected(self):
+    def is_connected(self) -> bool:
         """Return connected status."""
         return self._connected
 
     @property
-    def mac_address(self):
+    def mac_address(self) -> str:
         """Get MAC address."""
         return self._mac
 
     @property
-    def hostname(self):
+    def hostname(self) -> str:
         """Get Name."""
         return self._name
 
     @property
-    def ip_address(self):
+    def ip_address(self) -> str | None:
         """Get IP address."""
         return self._ip_address
 
     @property
-    def last_activity(self):
+    def last_activity(self) -> datetime | None:
         """Return device last activity."""
         return self._last_activity
 
@@ -274,7 +320,7 @@ class FritzBoxBaseEntity:
         return self._fritzbox_tools.mac
 
     @property
-    def device_info(self):
+    def device_info(self) -> DeviceInfo:
         """Return the device information."""
 
         return {

homeassistant/components/fritz/config_flow.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 
 import logging
 from typing import Any
-from urllib.parse import urlparse
+from urllib.parse import ParseResult, urlparse
 
 from fritzconnection.core.exceptions import FritzConnectionException, FritzSecurityError
 import voluptuous as vol
@@ -21,6 +21,7 @@ from homeassistant.config_entries import ConfigEntry, ConfigFlow, OptionsFlow
 from homeassistant.const import CONF_HOST, CONF_PASSWORD, CONF_PORT, CONF_USERNAME
 from homeassistant.core import callback
 from homeassistant.data_entry_flow import FlowResult
+from homeassistant.helpers.typing import DiscoveryInfoType
 
 from .common import FritzBoxTools
 from .const import (
@@ -42,23 +43,26 @@ class FritzBoxToolsFlowHandler(ConfigFlow, domain=DOMAIN):
 
     @staticmethod
     @callback
-    def async_get_options_flow(config_entry):
+    def async_get_options_flow(config_entry: ConfigEntry) -> OptionsFlow:
         """Get the options flow for this handler."""
         return FritzBoxToolsOptionsFlowHandler(config_entry)
 
-    def __init__(self):
+    def __init__(self) -> None:
         """Initialize FRITZ!Box Tools flow."""
-        self._host = None
-        self._entry = None
-        self._name = None
-        self._password = None
-        self._port = None
-        self._username = None
-        self.import_schema = None
-        self.fritz_tools = None
+        self._host: str | None = None
+        self._entry: ConfigEntry
+        self._name: str
+        self._password: str
+        self._port: int | None = None
+        self._username: str
+        self.fritz_tools: FritzBoxTools
 
-    async def fritz_tools_init(self):
+    async def fritz_tools_init(self) -> str | None:
         """Initialize FRITZ!Box Tools class."""
+
+        if not self._host or not self._port:
+            return None
+
         self.fritz_tools = FritzBoxTools(
             hass=self.hass,
             host=self._host,
@@ -87,7 +91,7 @@ class FritzBoxToolsFlowHandler(ConfigFlow, domain=DOMAIN):
         return None
 
     @callback
-    def _async_create_entry(self):
+    def _async_create_entry(self) -> FlowResult:
         """Async create flow handler entry."""
         return self.async_create_entry(
             title=self._name,
@@ -102,12 +106,14 @@ class FritzBoxToolsFlowHandler(ConfigFlow, domain=DOMAIN):
             },
         )
 
-    async def async_step_ssdp(self, discovery_info):
+    async def async_step_ssdp(self, discovery_info: DiscoveryInfoType) -> FlowResult:
         """Handle a flow initialized by discovery."""
-        ssdp_location = urlparse(discovery_info[ATTR_SSDP_LOCATION])
+        ssdp_location: ParseResult = urlparse(discovery_info[ATTR_SSDP_LOCATION])
         self._host = ssdp_location.hostname
         self._port = ssdp_location.port
-        self._name = discovery_info.get(ATTR_UPNP_FRIENDLY_NAME)
+        self._name = (
+            discovery_info.get(ATTR_UPNP_FRIENDLY_NAME) or self.fritz_tools.model
+        )
         self.context[CONF_HOST] = self._host
 
         if uuid := discovery_info.get(ATTR_UPNP_UDN):
@@ -130,7 +136,9 @@ class FritzBoxToolsFlowHandler(ConfigFlow, domain=DOMAIN):
         }
         return await self.async_step_confirm()
 
-    async def async_step_confirm(self, user_input=None):
+    async def async_step_confirm(
+        self, user_input: dict[str, Any] | None = None
+    ) -> FlowResult:
         """Handle user-confirmation of discovered node."""
         if user_input is None:
             return self._show_setup_form_confirm()
@@ -148,7 +156,7 @@ class FritzBoxToolsFlowHandler(ConfigFlow, domain=DOMAIN):
 
         return self._async_create_entry()
 
-    def _show_setup_form_init(self, errors=None):
+    def _show_setup_form_init(self, errors: dict[str, str] | None = None) -> FlowResult:
         """Show the setup form to the user."""
         return self.async_show_form(
             step_id="user",
@@ -163,7 +171,9 @@ class FritzBoxToolsFlowHandler(ConfigFlow, domain=DOMAIN):
             errors=errors or {},
         )
 
-    def _show_setup_form_confirm(self, errors=None):
+    def _show_setup_form_confirm(
+        self, errors: dict[str, str] | None = None
+    ) -> FlowResult:
         """Show the setup form to the user."""
         return self.async_show_form(
             step_id="confirm",
@@ -177,7 +187,9 @@ class FritzBoxToolsFlowHandler(ConfigFlow, domain=DOMAIN):
             errors=errors or {},
         )
 
-    async def async_step_user(self, user_input=None):
+    async def async_step_user(
+        self, user_input: dict[str, Any] | None = None
+    ) -> FlowResult:
         """Handle a flow initiated by the user."""
         if user_input is None:
             return self._show_setup_form_init()
@@ -197,24 +209,28 @@ class FritzBoxToolsFlowHandler(ConfigFlow, domain=DOMAIN):
 
         return self._async_create_entry()
 
-    async def async_step_reauth(self, data):
+    async def async_step_reauth(self, data: dict[str, Any]) -> FlowResult:
         """Handle flow upon an API authentication error."""
-        self._entry = self.hass.config_entries.async_get_entry(self.context["entry_id"])
+        if cfg_entry := self.hass.config_entries.async_get_entry(
+            self.context["entry_id"]
+        ):
+            self._entry = cfg_entry
         self._host = data[CONF_HOST]
         self._port = data[CONF_PORT]
         self._username = data[CONF_USERNAME]
         self._password = data[CONF_PASSWORD]
         return await self.async_step_reauth_confirm()
 
-    def _show_setup_form_reauth_confirm(self, user_input, errors=None):
+    def _show_setup_form_reauth_confirm(
+        self, user_input: dict[str, Any], errors: dict[str, str] | None = None
+    ) -> FlowResult:
         """Show the reauth form to the user."""
+        default_username = user_input.get(CONF_USERNAME)
         return self.async_show_form(
             step_id="reauth_confirm",
             data_schema=vol.Schema(
                 {
-                    vol.Required(
-                        CONF_USERNAME, default=user_input.get(CONF_USERNAME)
-                    ): str,
+                    vol.Required(CONF_USERNAME, default=default_username): str,
                     vol.Required(CONF_PASSWORD): str,
                 }
             ),
@@ -222,7 +238,9 @@ class FritzBoxToolsFlowHandler(ConfigFlow, domain=DOMAIN):
             errors=errors or {},
         )
 
-    async def async_step_reauth_confirm(self, user_input=None):
+    async def async_step_reauth_confirm(
+        self, user_input: dict[str, Any] | None = None
+    ) -> FlowResult:
         """Dialog that informs the user that reauth is required."""
         if user_input is None:
             return self._show_setup_form_reauth_confirm(
@@ -249,7 +267,7 @@ class FritzBoxToolsFlowHandler(ConfigFlow, domain=DOMAIN):
         await self.hass.config_entries.async_reload(self._entry.entry_id)
         return self.async_abort(reason="reauth_successful")
 
-    async def async_step_import(self, import_config):
+    async def async_step_import(self, import_config: dict[str, Any]) -> FlowResult:
         """Import a config entry from configuration.yaml."""
         return await self.async_step_user(
             {

homeassistant/components/fritz/device_tracker.py

@@ -1,6 +1,7 @@
 """Support for FRITZ!Box routers."""
 from __future__ import annotations
 
+import datetime
 import logging
 
 import voluptuous as vol
@@ -120,9 +121,9 @@ class FritzBoxTracker(ScannerEntity):
     def __init__(self, router: FritzBoxTools, device: FritzDevice) -> None:
         """Initialize a FRITZ!Box device."""
         self._router = router
-        self._mac = device.mac_address
-        self._name = device.hostname or DEFAULT_DEVICE_NAME
-        self._last_activity = device.last_activity
+        self._mac: str = device.mac_address
+        self._name: str = device.hostname or DEFAULT_DEVICE_NAME
+        self._last_activity: datetime.datetime | None = device.last_activity
         self._active = False
 
     @property

homeassistant/components/fritz/sensor.py

@@ -74,7 +74,10 @@ async def async_setup_entry(
     _LOGGER.debug("Setting up FRITZ!Box sensors")
     fritzbox_tools: FritzBoxTools = hass.data[DOMAIN][entry.entry_id]
 
-    if "WANIPConn1" not in fritzbox_tools.connection.services:
+    if (
+        not fritzbox_tools.connection
+        or "WANIPConn1" not in fritzbox_tools.connection.services
+    ):
         # Only routers are supported at the moment
         return