Improve tplink_lte typing (home-assistant#108393)

homeassistant/components/tplink_lte/__init__.py

@@ -1,6 +1,9 @@
 """Support for TP-Link LTE modems."""
+from __future__ import annotations
+
 import asyncio
 import logging
+from typing import Any
 
 import aiohttp
 import attr
@@ -15,7 +18,7 @@
     EVENT_HOMEASSISTANT_STOP,
     Platform,
 )
-from homeassistant.core import HomeAssistant, callback
+from homeassistant.core import Event, HomeAssistant, callback
 from homeassistant.helpers import config_validation as cv, discovery
 from homeassistant.helpers.aiohttp_client import async_create_clientsession
 from homeassistant.helpers.typing import ConfigType
@@ -59,20 +62,20 @@
 class ModemData:
     """Class for modem state."""
 
-    host = attr.ib()
-    modem = attr.ib()
+    host: str = attr.ib()
+    modem: tp_connected.Modem = attr.ib()
 
-    connected = attr.ib(init=False, default=True)
+    connected: bool = attr.ib(init=False, default=True)
 
 
 @attr.s
 class LTEData:
     """Shared state."""
 
-    websession = attr.ib()
+    websession: aiohttp.ClientSession = attr.ib()
     modem_data: dict[str, ModemData] = attr.ib(init=False, factory=dict)
 
-    def get_modem_data(self, config):
+    def get_modem_data(self, config: dict[str, Any]) -> ModemData | None:
         """Get the requested or the only modem_data value."""
         if CONF_HOST in config:
             return self.modem_data.get(config[CONF_HOST])
@@ -107,14 +110,16 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
     return True
 
 
-async def _setup_lte(hass, lte_config, delay=0):
+async def _setup_lte(
+    hass: HomeAssistant, lte_config: dict[str, Any], delay: int = 0
+) -> None:
     """Set up a TP-Link LTE modem."""
 
-    host = lte_config[CONF_HOST]
-    password = lte_config[CONF_PASSWORD]
+    host: str = lte_config[CONF_HOST]
+    password: str = lte_config[CONF_PASSWORD]
 
-    websession = hass.data[DATA_KEY].websession
-    modem = tp_connected.Modem(hostname=host, websession=websession)
+    lte_data: LTEData = hass.data[DATA_KEY]
+    modem = tp_connected.Modem(hostname=host, websession=lte_data.websession)
 
     modem_data = ModemData(host, modem)
 
@@ -124,28 +129,31 @@ async def _setup_lte(hass, lte_config, delay=0):
         retry_task = hass.loop.create_task(_retry_login(hass, modem_data, password))
 
         @callback
-        def cleanup_retry(event):
+        def cleanup_retry(event: Event) -> None:
             """Clean up retry task resources."""
             if not retry_task.done():
                 retry_task.cancel()
 
         hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, cleanup_retry)
 
 
-async def _login(hass, modem_data, password):
+async def _login(hass: HomeAssistant, modem_data: ModemData, password: str) -> None:
     """Log in and complete setup."""
     await modem_data.modem.login(password=password)
     modem_data.connected = True
-    hass.data[DATA_KEY].modem_data[modem_data.host] = modem_data
+    lte_data: LTEData = hass.data[DATA_KEY]
+    lte_data.modem_data[modem_data.host] = modem_data
 
-    async def cleanup(event):
+    async def cleanup(event: Event) -> None:
         """Clean up resources."""
         await modem_data.modem.logout()
 
     hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, cleanup)
 
 
-async def _retry_login(hass, modem_data, password):
+async def _retry_login(
+    hass: HomeAssistant, modem_data: ModemData, password: str
+) -> None:
     """Sleep and retry setup."""
 
     _LOGGER.warning("Could not connect to %s. Will keep trying", modem_data.host)

homeassistant/components/tplink_lte/notify.py

@@ -2,6 +2,7 @@
 from __future__ import annotations
 
 import logging
+from typing import Any
 
 import attr
 import tp_connected
@@ -11,7 +12,7 @@
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType
 
-from . import DATA_KEY
+from . import DATA_KEY, LTEData
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -31,13 +32,14 @@ async def async_get_service(
 class TplinkNotifyService(BaseNotificationService):
     """Implementation of a notification service."""
 
-    hass = attr.ib()
-    config = attr.ib()
+    hass: HomeAssistant = attr.ib()
+    config: dict[str, Any] = attr.ib()
 
-    async def async_send_message(self, message="", **kwargs):
+    async def async_send_message(self, message: str = "", **kwargs: Any) -> None:
         """Send a message to a user."""
 
-        modem_data = self.hass.data[DATA_KEY].get_modem_data(self.config)
+        lte_data: LTEData = self.hass.data[DATA_KEY]
+        modem_data = lte_data.get_modem_data(self.config)
         if not modem_data:
             _LOGGER.error("No modem available")
             return