MyPy

Author: andrew-codechimp

custom_components/battery_notes/__init__.py

@@ -147,6 +147,7 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: BatteryNotesConfi
     """Set up a config entry."""
 
     data = hass.data[MY_KEY]
+    assert(data.store)
     config_entry.runtime_data = BatteryNotesData(
         domain_config=data,
         store=data.store,
@@ -174,9 +175,7 @@ async def async_remove_entry(hass: HomeAssistant, config_entry: BatteryNotesConf
     # Remove any issues raised
     ir.async_delete_issue(hass, DOMAIN, f"missing_device_{config_entry.entry_id}")
 
-    config_entry.runtime_data.coordinator.async_update_device_config
-
-    if not config_entry.runtime_data.coordinator.device_id:
+    if not config_entry.runtime_data.coordinator or not config_entry.runtime_data.coordinator.device_id:
         return
 
     data = {ATTR_REMOVE: True}
@@ -201,7 +200,7 @@ async def async_remove_entry(hass: HomeAssistant, config_entry: BatteryNotesConf
 
     if wrapped_battery_entity_entry.hidden_by == er.RegistryEntryHider.INTEGRATION:
         entity_registry.async_update_entity(
-            config_entry.runtime_data.wrapped_battery.entity_id, hidden_by=None
+            config_entry.runtime_data.coordinator.wrapped_battery.entity_id, hidden_by=None
         )
         _LOGGER.debug(
             "Unhidden Original Battery for device%s", config_entry.runtime_data.coordinator.device_id

custom_components/battery_notes/binary_sensor.py

@@ -162,6 +162,7 @@ async def async_registry_updated(
             )
 
     coordinator = config_entry.runtime_data.coordinator
+    assert(coordinator)
 
     config_entry.async_on_unload(
         async_track_entity_registry_updated_event(

custom_components/battery_notes/button.py

@@ -128,6 +128,7 @@ async def async_registry_updated(event: Event[er.EventEntityRegistryUpdatedData]
             )
 
     coordinator = config_entry.runtime_data.coordinator
+    assert(coordinator)
 
     config_entry.async_on_unload(
         async_track_entity_registry_updated_event(

custom_components/battery_notes/config_flow.py

@@ -185,7 +185,8 @@ async def async_step_device(
                     device_entry.hw_version,
                 )
 
-                library = await Library.factory(self.hass)
+                library = Library(self.hass)
+                await library.load_libraries()
 
                 # Set defaults if not found in library
                 self.data[CONF_BATTERY_QUANTITY] = 1
@@ -284,7 +285,8 @@ async def async_step_entity(
                             device_entry.hw_version,
                         )
 
-                        library = await Library.factory(self.hass)
+                        library = Library(self.hass)
+                        await library.load_libraries()
 
                         device_battery_details = (
                             await library.get_device_battery_details(self.model_info)

custom_components/battery_notes/coordinator.py

@@ -96,6 +96,7 @@ class BatteryNotesData:
 class BatteryNotesCoordinator(DataUpdateCoordinator[None]):
     """Define an object to hold Battery Notes device."""
 
+    config_entry: BatteryNotesConfigEntry
     device_id: str | None = None
     source_entity_id: str | None = None
     device_name: str
@@ -176,7 +177,7 @@ def __init__(
                     last_replaced = device_entry.created_at.strftime(
                         "%Y-%m-%dT%H:%M:%S:%f"
                     )
-            else:
+            elif self.source_entity_id:
                 entity = entity_registry.async_get(self.source_entity_id)
                 if entity and entity.created_at.year > 1970:
                     last_replaced = entity.created_at.strftime("%Y-%m-%dT%H:%M:%S:%f")
@@ -295,7 +296,9 @@ def _link_device(self) -> bool:
                     if self.wrapped_battery:
                         break
 
-            device_entry = device_registry.async_get(self.device_id)
+            device_entry = None
+            if self.device_id:
+                device_entry = device_registry.async_get(self.device_id)
             if device_entry:
                 self.device_name = (
                     device_entry.name_by_user or device_entry.name or self.config_entry.title

custom_components/battery_notes/discovery.py

@@ -10,7 +10,6 @@
 from homeassistant.const import CONF_DEVICE_ID
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers import discovery_flow
-from homeassistant.helpers.typing import ConfigType
 
 from .common import get_device_model_id
 from .const import (
@@ -22,6 +21,7 @@
     CONF_MODEL_ID,
     DOMAIN,
 )
+from .coordinator import BatteryNotesDomainConfig
 from .library import DeviceBatteryDetails, Library, ModelInfo
 
 _LOGGER = logging.getLogger(__name__)
@@ -77,7 +77,7 @@ class DiscoveryManager:
     so the user can add them to their HA instance.
     """
 
-    def __init__(self, hass: HomeAssistant, ha_config: ConfigType) -> None:
+    def __init__(self, hass: HomeAssistant, ha_config: BatteryNotesDomainConfig) -> None:
         """Init."""
         self.hass = hass
         self.ha_config = ha_config

custom_components/battery_notes/library.py

@@ -101,18 +101,18 @@ def _load_library_json(library_file: str) -> dict[str, Any]:
                 json_default_path,
             )
 
-    @staticmethod
-    async def factory(hass: HomeAssistant) -> Library:
-        """Return the library or create."""
-
-        domain_config = hass.data[MY_KEY]
-        if domain_config.library:
-            return domain_config.library
-
-        library = Library(hass)
-        await library.load_libraries()
-        domain_config.library = library
-        return library
+    # @staticmethod
+    # async def factory(hass: HomeAssistant) -> Library:
+    #     """Return the library or create."""
+
+    #     domain_config = hass.data[MY_KEY]
+    #     if domain_config.library:
+    #         return domain_config.library
+
+    #     library = Library(hass)
+    #     await library.load_libraries()
+    #     domain_config.library = library
+    #     return library
 
     async def get_device_battery_details(
         self,

custom_components/battery_notes/sensor.py

@@ -164,6 +164,7 @@ async def async_registry_updated(event: Event[er.EventEntityRegistryUpdatedData]
     )
 
     coordinator = config_entry.runtime_data.coordinator
+    assert(coordinator)
 
     if not coordinator.fake_device:
         device_id = async_add_to_device(hass, config_entry)

custom_components/battery_notes/services.py

@@ -77,7 +77,8 @@ async def handle_battery_replaced(call: ServiceCall) -> ServiceResponse:
             config_entry: BatteryNotesConfigEntry
             for config_entry in hass.config_entries.async_loaded_entries(DOMAIN):
                 coordinator = config_entry.runtime_data.coordinator
-                if coordinator.source_entity_id == source_entity_id:
+                assert(coordinator)
+                if coordinator.source_entity_id and coordinator.source_entity_id == source_entity_id:
 
                     coordinator.last_replaced =datetime_replaced
                     await coordinator.async_request_refresh()
@@ -120,9 +121,9 @@ async def handle_battery_replaced(call: ServiceCall) -> ServiceResponse:
                 )
                 return None
 
-            config_entry: BatteryNotesConfigEntry
             for config_entry in hass.config_entries.async_loaded_entries(DOMAIN):
                 coordinator = config_entry.runtime_data.coordinator
+                assert(coordinator)
                 if coordinator.device_id == device_id:
                     coordinator.last_replaced = datetime_replaced
                     await coordinator.async_request_refresh()
@@ -168,6 +169,7 @@ async def handle_battery_last_reported(call: ServiceCall) -> ServiceResponse:
         config_entry: BatteryNotesConfigEntry
         for config_entry in hass.config_entries.async_loaded_entries(DOMAIN):
             coordinator = config_entry.runtime_data.coordinator
+            assert(coordinator)
 
             if coordinator.wrapped_battery and coordinator.last_reported:
                 time_since_lastreported = (
@@ -206,6 +208,7 @@ async def handle_battery_low(call: ServiceCall) -> ServiceResponse:
         config_entry: BatteryNotesConfigEntry
         for config_entry in hass.config_entries.async_loaded_entries(DOMAIN):
             coordinator = config_entry.runtime_data.coordinator
+            assert(coordinator)
 
             if coordinator.battery_low is True:
                 hass.bus.async_fire(