Strict typing for homekit part 1 (#67657)

Co-authored-by: Franck Nijhof <frenck@frenck.nl>

homeassistant/components/homekit/aidmanager.py

@@ -9,13 +9,15 @@ can't change the hash without causing breakages for HA users.
 
 This module generates and stores them in a HA storage.
 """
+from __future__ import annotations
+
+from collections.abc import Generator
 import random
 
 from fnvhash import fnv1a_32
 
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
-from homeassistant.helpers.entity_registry import RegistryEntry
+from homeassistant.helpers.entity_registry import EntityRegistry, RegistryEntry
 from homeassistant.helpers.storage import Store
 
 from .util import get_aid_storage_filename_for_entry_id
@@ -32,12 +34,12 @@ AID_MIN = 2
 AID_MAX = 18446744073709551615
 
 
-def get_system_unique_id(entity: RegistryEntry):
+def get_system_unique_id(entity: RegistryEntry) -> str:
     """Determine the system wide unique_id for an entity."""
     return f"{entity.platform}.{entity.domain}.{entity.unique_id}"
 
 
-def _generate_aids(unique_id: str, entity_id: str) -> int:
+def _generate_aids(unique_id: str | None, entity_id: str) -> Generator[int, None, None]:
     """Generate accessory aid."""
 
     if unique_id:
@@ -65,39 +67,41 @@ class AccessoryAidStorage:
     persist over reboots.
     """
 
-    def __init__(self, hass: HomeAssistant, entry: ConfigEntry) -> None:
+    def __init__(self, hass: HomeAssistant, entry_id: str) -> None:
         """Create a new entity map store."""
         self.hass = hass
-        self.allocations = {}
-        self.allocated_aids = set()
-        self._entry = entry
-        self.store = None
-        self._entity_registry = None
+        self.allocations: dict[str, int] = {}
+        self.allocated_aids: set[int] = set()
+        self._entry_id = entry_id
+        self.store: Store | None = None
+        self._entity_registry: EntityRegistry | None = None
 
-    async def async_initialize(self):
+    async def async_initialize(self) -> None:
         """Load the latest AID data."""
         self._entity_registry = (
             await self.hass.helpers.entity_registry.async_get_registry()
         )
-        aidstore = get_aid_storage_filename_for_entry_id(self._entry)
+        aidstore = get_aid_storage_filename_for_entry_id(self._entry_id)
         self.store = Store(self.hass, AID_MANAGER_STORAGE_VERSION, aidstore)
 
         if not (raw_storage := await self.store.async_load()):
             # There is no data about aid allocations yet
             return
 
+        assert isinstance(raw_storage, dict)
         self.allocations = raw_storage.get(ALLOCATIONS_KEY, {})
         self.allocated_aids = set(self.allocations.values())
 
-    def get_or_allocate_aid_for_entity_id(self, entity_id: str):
+    def get_or_allocate_aid_for_entity_id(self, entity_id: str) -> int:
         """Generate a stable aid for an entity id."""
+        assert self._entity_registry is not None
         if not (entity := self._entity_registry.async_get(entity_id)):
             return self.get_or_allocate_aid(None, entity_id)
 
         sys_unique_id = get_system_unique_id(entity)
         return self.get_or_allocate_aid(sys_unique_id, entity_id)
 
-    def get_or_allocate_aid(self, unique_id: str, entity_id: str):
+    def get_or_allocate_aid(self, unique_id: str | None, entity_id: str) -> int:
         """Allocate (and return) a new aid for an accessory."""
         if unique_id and unique_id in self.allocations:
             return self.allocations[unique_id]
@@ -119,7 +123,7 @@ class AccessoryAidStorage:
             f"Unable to generate unique aid allocation for {entity_id} [{unique_id}]"
         )
 
-    def delete_aid(self, storage_key: str):
+    def delete_aid(self, storage_key: str) -> None:
         """Delete an aid allocation."""
         if storage_key not in self.allocations:
             return
@@ -129,15 +133,17 @@ class AccessoryAidStorage:
         self.async_schedule_save()
 
     @callback
-    def async_schedule_save(self):
+    def async_schedule_save(self) -> None:
         """Schedule saving the entity map cache."""
+        assert self.store is not None
         self.store.async_delay_save(self._data_to_save, AID_MANAGER_SAVE_DELAY)
 
-    async def async_save(self):
+    async def async_save(self) -> None:
         """Save the entity map cache."""
+        assert self.store is not None
         return await self.store.async_save(self._data_to_save())
 
     @callback
-    def _data_to_save(self):
+    def _data_to_save(self) -> dict:
         """Return data of entity map to store in a file."""
         return {ALLOCATIONS_KEY: self.allocations}