Improve typing of Tasmota (3/3) (#52748)

homeassistant/components/tasmota/mixins.py

@@ -4,6 +4,13 @@ from __future__ import annotations
 import logging
 from typing import Any
 
+from hatasmota.entity import (
+    TasmotaAvailability as HATasmotaAvailability,
+    TasmotaEntity as HATasmotaEntity,
+    TasmotaEntityConfig,
+)
+from hatasmota.models import DiscoveryHashType
+
 from homeassistant.components.mqtt import (
     async_subscribe_connection_status,
     is_connected as mqtt_connected,
@@ -11,7 +18,7 @@ from homeassistant.components.mqtt import (
 from homeassistant.core import callback
 from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
-from homeassistant.helpers.entity import Entity
+from homeassistant.helpers.entity import DeviceInfo, Entity
 
 from .discovery import (
     TASMOTA_DISCOVERY_ENTITY_UPDATED,
@@ -25,48 +32,50 @@ _LOGGER = logging.getLogger(__name__)
 class TasmotaEntity(Entity):
     """Base class for Tasmota entities."""
 
-    def __init__(self, tasmota_entity) -> None:
+    def __init__(self, tasmota_entity: HATasmotaEntity) -> None:
         """Initialize."""
         self._tasmota_entity = tasmota_entity
         self._unique_id = tasmota_entity.unique_id
 
-    async def async_added_to_hass(self):
+    async def async_added_to_hass(self) -> None:
         """Subscribe to MQTT events."""
         await self._subscribe_topics()
 
-    async def async_will_remove_from_hass(self):
+    async def async_will_remove_from_hass(self) -> None:
         """Unsubscribe when removed."""
         await self._tasmota_entity.unsubscribe_topics()
         await super().async_will_remove_from_hass()
 
-    async def discovery_update(self, update, write_state=True):
+    async def discovery_update(
+        self, update: TasmotaEntityConfig, write_state: bool = True
+    ) -> None:
         """Handle updated discovery message."""
         self._tasmota_entity.config_update(update)
         await self._subscribe_topics()
         if write_state:
             self.async_write_ha_state()
 
-    async def _subscribe_topics(self):
+    async def _subscribe_topics(self) -> None:
         """(Re)Subscribe to topics."""
         await self._tasmota_entity.subscribe_topics()
 
     @property
-    def device_info(self):
+    def device_info(self) -> DeviceInfo:
         """Return a device description for device registry."""
         return {"connections": {(CONNECTION_NETWORK_MAC, self._tasmota_entity.mac)}}
 
     @property
-    def name(self):
+    def name(self) -> str | None:
         """Return the name of the binary sensor."""
         return self._tasmota_entity.name
 
     @property
-    def should_poll(self):
+    def should_poll(self) -> bool:
         """Return the polling state."""
         return False
 
     @property
-    def unique_id(self):
+    def unique_id(self) -> str:
         """Return a unique ID."""
         return self._unique_id
 
@@ -99,7 +108,9 @@ class TasmotaOnOffEntity(TasmotaEntity):
 class TasmotaAvailability(TasmotaEntity):
     """Mixin used for platforms that report availability."""
 
-    def __init__(self, **kwds) -> None:
+    _tasmota_entity: HATasmotaAvailability
+
+    def __init__(self, **kwds: Any) -> None:
         """Initialize the availability mixin."""
         self._available = False
         super().__init__(**kwds)
@@ -120,7 +131,7 @@ class TasmotaAvailability(TasmotaEntity):
         self.async_write_ha_state()
 
     @callback
-    def async_mqtt_connected(self, _):
+    def async_mqtt_connected(self, _: bool) -> None:
         """Update state on connection/disconnection to MQTT broker."""
         if not self.hass.is_stopping:
             if not mqtt_connected(self.hass):
@@ -136,7 +147,7 @@ class TasmotaAvailability(TasmotaEntity):
 class TasmotaDiscoveryUpdate(TasmotaEntity):
     """Mixin used to handle updated discovery message."""
 
-    def __init__(self, discovery_hash, **kwds) -> None:
+    def __init__(self, discovery_hash: DiscoveryHashType, **kwds: Any) -> None:
         """Initialize the discovery update mixin."""
         self._discovery_hash = discovery_hash
         self._removed_from_hass = False
@@ -147,7 +158,7 @@ class TasmotaDiscoveryUpdate(TasmotaEntity):
         self._removed_from_hass = False
         await super().async_added_to_hass()
 
-        async def discovery_callback(config):
+        async def discovery_callback(config: TasmotaEntityConfig) -> None:
             """Handle discovery update."""
             _LOGGER.debug(
                 "Got update for entity with hash: %s '%s'",