diff --git a/custom_components/foxess_modbus/__init__.py b/custom_components/foxess_modbus/__init__.py index 696cf35..24a1ff8 100755 --- a/custom_components/foxess_modbus/__init__.py +++ b/custom_components/foxess_modbus/__init__.py @@ -18,10 +18,11 @@ 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 @@ -29,7 +30,6 @@ 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 diff --git a/custom_components/foxess_modbus/binary_sensor.py b/custom_components/foxess_modbus/binary_sensor.py index 6cdd21c..6f2988a 100755 --- a/custom_components/foxess_modbus/binary_sensor.py +++ b/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)) diff --git a/custom_components/foxess_modbus/common/types.py b/custom_components/foxess_modbus/common/types.py index 3e13548..e713f9e 100644 --- a/custom_components/foxess_modbus/common/types.py +++ b/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] diff --git a/custom_components/foxess_modbus/const.py b/custom_components/foxess_modbus/const.py index 0966c2c..916b9fd 100755 --- a/custom_components/foxess_modbus/const.py +++ b/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" diff --git a/custom_components/foxess_modbus/number.py b/custom_components/foxess_modbus/number.py index 0e91b6b..47ba8c6 100644 --- a/custom_components/foxess_modbus/number.py +++ b/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)) diff --git a/custom_components/foxess_modbus/select.py b/custom_components/foxess_modbus/select.py index 18097b5..b0a49bf 100644 --- a/custom_components/foxess_modbus/select.py +++ b/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)) diff --git a/custom_components/foxess_modbus/sensor.py b/custom_components/foxess_modbus/sensor.py index 8b0a543..238efdc 100755 --- a/custom_components/foxess_modbus/sensor.py +++ b/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)]) diff --git a/custom_components/foxess_modbus/services/websocket_api.py b/custom_components/foxess_modbus/services/websocket_api.py index 0a057c9..6768611 100644 --- a/custom_components/foxess_modbus/services/websocket_api.py +++ b/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(