Merge pull request #579 from canton7/bugfix/charge-period

Fix other things which depended on the structure of hass.data

custom_components/foxess_modbus/__init__.py

@@ -18,18 +18,18 @@
 from slugify import slugify
 
 from .client.modbus_client import ModbusClient
+from .common.types import HassData
+from .common.types import HassDataEntry
 from .const import ADAPTER_ID
 from .const import ADAPTER_WAS_MIGRATED
 from .const import CONFIG_SAVE_TIME
-from .const import CONTROLLERS
 from .const import DOMAIN
 from .const import ENTITY_ID_PREFIX
 from .const import FRIENDLY_NAME
 from .const import HOST
 from .const import INVERTER_CONN
 from .const import INVERTERS
 from .const import MAX_READ
-from .const import MODBUS_CLIENTS
 from .const import MODBUS_SLAVE
 from .const import MODBUS_TYPE
 from .const import PLATFORMS
@@ -64,7 +64,9 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     entry_options = copy.deepcopy(dict(entry.options))
 
     # Create this before throwing ConfigEntryAuthFailed, so the sensors, etc, platforms don't fail
-    hass.data.setdefault(DOMAIN, {}).setdefault(entry.entry_id, {})[INVERTERS] = []
+    hass.data.setdefault(DOMAIN, HassData()).setdefault(
+        entry.entry_id, HassDataEntry(controllers=[], modbus_clients=[])
+    )
 
     for platform in PLATFORMS:
         if entry_options.get(platform, True):
@@ -122,9 +124,10 @@ def create_controller(client: ModbusClient, inverter: dict[str, Any]) -> None:
     update_charge_period_service.register(hass, controllers)
     websocket_api.register(hass)
 
-    hass.data[DOMAIN][entry.entry_id][CONTROLLERS] = controllers
-    hass.data[DOMAIN][entry.entry_id][MODBUS_CLIENTS] = clients.values()
-    hass.data[DOMAIN][entry.entry_id]["unload"] = entry.add_update_listener(async_reload_entry)
+    hass_data: HassData = hass.data[DOMAIN]
+    hass_data[entry.entry_id]["controllers"] = controllers
+    hass_data[entry.entry_id]["modbus_clients"] = list(clients.values())
+    hass_data[entry.entry_id]["unload"] = entry.add_update_listener(async_reload_entry)
 
     return True
 
@@ -249,14 +252,16 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     )
 
     if unloaded:
-        controllers = hass.data[DOMAIN][entry.entry_id][INVERTERS]
-        for _, controller in controllers:
+        hass_data: HassData = hass.data[DOMAIN]
+        controllers = hass_data[entry.entry_id]["controllers"]
+        for controller in controllers:
             controller.unload()
-        clients = hass.data[DOMAIN][entry.entry_id][MODBUS_CLIENTS]
+
+        clients = hass_data[entry.entry_id]["modbus_clients"]
         await asyncio.gather(*[client.close() for client in clients])
 
-        hass.data[DOMAIN][entry.entry_id]["unload"]()
-        hass.data[DOMAIN].pop(entry.entry_id)
+        hass_data[entry.entry_id]["unload"]()
+        hass_data.pop(entry.entry_id)
 
     return unloaded
 

custom_components/foxess_modbus/binary_sensor.py

@@ -7,7 +7,7 @@
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import CONTROLLERS
+from .common.types import HassData
 from .const import DOMAIN
 from .inverter_profiles import create_entities
 
@@ -17,7 +17,8 @@
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry, async_add_devices: AddEntitiesCallback) -> None:
     """Setup sensor platform."""
 
-    controllers = hass.data[DOMAIN][entry.entry_id][CONTROLLERS]
+    hass_data: HassData = hass.data[DOMAIN]
+    controllers = hass_data[entry.entry_id]["controllers"]
 
     for controller in controllers:
         async_add_devices(create_entities(BinarySensorEntity, controller))

custom_components/foxess_modbus/common/types.py

@@ -5,6 +5,15 @@
 from enum import IntEnum
 from enum import StrEnum
 from enum import auto
+from typing import TYPE_CHECKING
+from typing import Callable
+from typing import NotRequired
+from typing import TypeAlias
+from typing import TypedDict
+
+if TYPE_CHECKING:
+    from ..client.modbus_client import ModbusClient
+    from ..modbus_controller import ModbusController
 
 
 class RegisterType(Enum):
@@ -69,3 +78,12 @@ class RegisterPollType(IntEnum):
     # These must be ordered from least frequent to most frequent
     ON_CONNECTION = 0
     PERIODICALLY = 1
+
+
+class HassDataEntry(TypedDict):
+    controllers: list["ModbusController"]
+    modbus_clients: list["ModbusClient"]
+    unload: NotRequired[Callable[[], None]]
+
+
+HassData: TypeAlias = dict[str, HassDataEntry]

custom_components/foxess_modbus/const.py

@@ -48,8 +48,6 @@
 INVERTER_BASE = "inverter_base"
 INVERTER_CONN = "inverter_conn"
 INVERTERS = "inverters"
-CONTROLLERS = "controllers"
-MODBUS_CLIENTS = "modbus_clients"
 
 CONFIG_SAVE_TIME = "save_time"
 

custom_components/foxess_modbus/number.py

@@ -7,7 +7,7 @@
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import CONTROLLERS
+from .common.types import HassData
 from .const import DOMAIN
 from .inverter_profiles import create_entities
 
@@ -17,7 +17,8 @@
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry, async_add_devices: AddEntitiesCallback) -> None:
     """Setup numbers platform."""
 
-    controllers = hass.data[DOMAIN][entry.entry_id][CONTROLLERS]
+    hass_data: HassData = hass.data[DOMAIN]
+    controllers = hass_data[entry.entry_id]["controllers"]
 
     for controller in controllers:
         async_add_devices(create_entities(NumberEntity, controller))

custom_components/foxess_modbus/select.py

@@ -7,7 +7,7 @@
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import CONTROLLERS
+from .common.types import HassData
 from .const import DOMAIN
 from .inverter_profiles import create_entities
 
@@ -17,7 +17,8 @@
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry, async_add_devices: AddEntitiesCallback) -> None:
     """Setup select platform."""
 
-    controllers = hass.data[DOMAIN][entry.entry_id][CONTROLLERS]
+    hass_data: HassData = hass.data[DOMAIN]
+    controllers = hass_data[entry.entry_id]["controllers"]
 
     for controller in controllers:
         async_add_devices(create_entities(SelectEntity, controller))

custom_components/foxess_modbus/sensor.py

@@ -7,7 +7,7 @@
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import CONTROLLERS
+from .common.types import HassData
 from .const import DOMAIN
 from .entities.connection_status_sensor import ConnectionStatusSensor
 from .inverter_profiles import create_entities
@@ -18,7 +18,8 @@
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry, async_add_devices: AddEntitiesCallback) -> None:
     """Setup sensor platform."""
 
-    controllers = hass.data[DOMAIN][entry.entry_id][CONTROLLERS]
+    hass_data: HassData = hass.data[DOMAIN]
+    controllers = hass_data[entry.entry_id]["controllers"]
 
     for controller in controllers:
         async_add_devices([ConnectionStatusSensor(controller)])

custom_components/foxess_modbus/services/websocket_api.py

@@ -6,9 +6,9 @@
 from homeassistant.core import callback
 from homeassistant.helpers import config_validation as cv
 
+from ..common.types import HassData
 from ..const import DOMAIN
 from ..const import FRIENDLY_NAME
-from ..const import INVERTERS
 from .utils import get_controller_from_friendly_name_or_device_id
 
 
@@ -24,8 +24,9 @@ def register(hass: HomeAssistant) -> None:
 )
 @callback
 def get_charge_periods(hass: HomeAssistant, connection: websocket_api.ActiveConnection, msg: dict[str, Any]) -> None:
-    inverter_controllers = [x for entry in hass.data[DOMAIN].values() for x in entry[INVERTERS]]
-    controller = get_controller_from_friendly_name_or_device_id(msg["inverter"], inverter_controllers, hass)
+    hass_data: HassData = hass.data[DOMAIN]
+    controllers = [x for entry in hass_data.values() for x in entry["controllers"]]
+    controller = get_controller_from_friendly_name_or_device_id(msg["inverter"], controllers, hass)
     charge_periods = []
     for charge_period in controller.charge_periods:
         charge_periods.append(