Use entry.runtime_data (#82)

custom_components/amshan/__init__.py

@@ -22,14 +22,20 @@
     CONF_CONNECTION_TYPE,
     CONF_MQTT_TOPICS,
     CONF_TCP_HOST,
-    DOMAIN,
 )
 from .metercon import async_setup_meter_mqtt_subscriptions, setup_meter_connection
 
 _LOGGER: logging.Logger = logging.getLogger(__name__)
 
 PLATFORM_TYPE = ha_const.Platform.SENSOR
 
+type AmsHanConfigEntry = ConfigEntry[AmsHanData]
+
+
+@dataclass
+class AmsHanData:
+    integration: AmsHanIntegration
+
 
 class ConnectionType(Enum):
     """Meter connection type."""
@@ -108,11 +114,10 @@ def stop_receive(self) -> None:
 
 async def async_setup(hass: HomeAssistant, _: ConfigType) -> bool:
     """Set up the amshan component."""
-    hass.data[DOMAIN] = {}
     return True
 
 
-async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool:
+async def async_setup_entry(hass: HomeAssistant, config_entry: AmsHanConfigEntry) -> bool:
     """Set up amshan from a config entry."""
     integration = AmsHanIntegration()
 
@@ -133,7 +138,8 @@ async def on_hass_stop(event: Event) -> None:
         config_entry.add_update_listener(async_config_entry_changed)
     )
 
-    hass.data[DOMAIN][config_entry.entry_id] = integration
+    config_entry.runtime_data = AmsHanData(integration)
+
     await hass.config_entries.async_forward_entry_setup(config_entry, PLATFORM_TYPE)
 
     _LOGGER.debug("async_setup_entry complete.")
@@ -184,7 +190,7 @@ async def async_migrate_config_entry(
 
 
 async def async_unload_entry(
-    hass: HomeAssistant, config_entry: ConfigEntry
+    hass: HomeAssistant, config_entry: AmsHanConfigEntry
 ) -> bool:
     """Handle removal of an entry."""
     is_plaform_unload_success = await hass.config_entries.async_forward_entry_unload(
@@ -193,8 +199,7 @@ async def async_unload_entry(
 
     if is_plaform_unload_success:
         _LOGGER.info("Integrations is unloading.")
-        ctx: AmsHanIntegration = hass.data[DOMAIN].pop(config_entry.entry_id)
-        await ctx.async_close_all()
+        await config_entry.runtime_data.integration.async_close_all()
 
     return is_plaform_unload_success
 

custom_components/amshan/config_flow.py

@@ -5,7 +5,7 @@
 import logging
 import os
 import socket
-from typing import Any, cast
+from typing import Any
 
 import async_timeout
 from han import autodecoder, common as han_type, obis_map

custom_components/amshan/sensor.py

@@ -23,7 +23,7 @@
 from homeassistant.helpers import dispatcher, entity, restore_state
 from homeassistant.util import dt as dt_util
 
-from . import AmsHanIntegration, MeterInfo, StopMessage
+from . import AmsHanConfigEntry, MeterInfo, StopMessage
 from .const import (
     CONF_OPTIONS_SCALE_FACTOR,
     DOMAIN,
@@ -238,19 +238,18 @@ class AmsHanSensorEntityDescription(SensorEntityDescription):
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: AmsHanConfigEntry,
     async_add_entities: Callable[[list[entity.Entity], bool], None],
 ):
     """Add hantest sensor platform from a config_entry."""
     _LOGGER.debug("Sensor async_setup_entry starting.")
 
-    integration: AmsHanIntegration = hass.data[DOMAIN][config_entry.entry_id]
     processor: MeterMeasureProcessor = MeterMeasureProcessor(
-        hass, config_entry, async_add_entities, integration.measure_queue
+        hass, config_entry, async_add_entities, config_entry.runtime_data.integration.measure_queue
     )
 
     # start processing loop task
-    integration.add_task(hass.loop.create_task(processor.async_process_measures_loop()))
+    config_entry.runtime_data.integration.add_task(hass.loop.create_task(processor.async_process_measures_loop()))
 
     _LOGGER.debug("Sensor async_setup_entry ended.")