From edb06c58fa5bb523d877135cdcf33f1d78713fb9 Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Sat, 11 Mar 2023 18:04:19 -1000 Subject: [PATCH] Add some more typing to screenlogic (#88522) --- .../components/screenlogic/__init__.py | 29 ++++++------ .../components/screenlogic/entity.py | 45 ++++++++++++------- 2 files changed, 45 insertions(+), 29 deletions(-) diff --git a/homeassistant/components/screenlogic/__init__.py b/homeassistant/components/screenlogic/__init__.py index ad2f9c64f3e..a0accef0799 100644 --- a/homeassistant/components/screenlogic/__init__.py +++ b/homeassistant/components/screenlogic/__init__.py @@ -90,23 +90,23 @@ async def async_update_listener(hass: HomeAssistant, entry: ConfigEntry) -> None await hass.config_entries.async_reload(entry.entry_id) -async def async_get_connect_info(hass: HomeAssistant, entry: ConfigEntry): +async def async_get_connect_info( + hass: HomeAssistant, entry: ConfigEntry +) -> dict[str, str | int]: """Construct connect_info from configuration entry and returns it to caller.""" mac = entry.unique_id # Attempt to rediscover gateway to follow IP changes discovered_gateways = await async_discover_gateways_by_unique_id(hass) if mac in discovered_gateways: - connect_info = discovered_gateways[mac] - else: - _LOGGER.warning("Gateway rediscovery failed") - # Static connection defined or fallback from discovery - connect_info = { - SL_GATEWAY_NAME: name_for_mac(mac), - SL_GATEWAY_IP: entry.data[CONF_IP_ADDRESS], - SL_GATEWAY_PORT: entry.data[CONF_PORT], - } + return discovered_gateways[mac] - return connect_info + _LOGGER.warning("Gateway rediscovery failed") + # Static connection defined or fallback from discovery + return { + SL_GATEWAY_NAME: name_for_mac(mac), + SL_GATEWAY_IP: entry.data[CONF_IP_ADDRESS], + SL_GATEWAY_PORT: entry.data[CONF_PORT], + } class ScreenlogicDataUpdateCoordinator(DataUpdateCoordinator): @@ -143,7 +143,7 @@ class ScreenlogicDataUpdateCoordinator(DataUpdateCoordinator): """Return the gateway data.""" return self.gateway.get_data() - async def _async_update_configured_data(self): + async def _async_update_configured_data(self) -> None: """Update data sets based on equipment config.""" equipment_flags = self.gateway.get_data()[SL_DATA.KEY_CONFIG]["equipment_flags"] if not self.gateway.is_client: @@ -155,7 +155,7 @@ class ScreenlogicDataUpdateCoordinator(DataUpdateCoordinator): if equipment_flags & EQUIPMENT.FLAG_CHLORINATOR: await self.gateway.async_get_scg() - async def _async_update_data(self): + async def _async_update_data(self) -> None: """Fetch data from the Screenlogic gateway.""" try: await self._async_update_configured_data() @@ -165,8 +165,9 @@ class ScreenlogicDataUpdateCoordinator(DataUpdateCoordinator): return None - async def _async_reconnect_update_data(self): + async def _async_reconnect_update_data(self) -> None: """Attempt to reconnect to the gateway and fetch data.""" + assert self.config_entry is not None try: # Clean up the previous connection as we're about to create a new one await self.gateway.async_disconnect() diff --git a/homeassistant/components/screenlogic/entity.py b/homeassistant/components/screenlogic/entity.py index 80b4df4d2de..4ea23395c5a 100644 --- a/homeassistant/components/screenlogic/entity.py +++ b/homeassistant/components/screenlogic/entity.py @@ -1,9 +1,10 @@ """Base ScreenLogicEntity definitions.""" +from datetime import datetime import logging from typing import Any -# from screenlogicpy import ScreenLogicError, ScreenLogicGateway -from screenlogicpy.const import DATA as SL_DATA, EQUIPMENT, ON_OFF +from screenlogicpy import ScreenLogicGateway +from screenlogicpy.const import CODE, DATA as SL_DATA, EQUIPMENT, ON_OFF from homeassistant.core import callback from homeassistant.exceptions import HomeAssistantError @@ -19,7 +20,12 @@ _LOGGER = logging.getLogger(__name__) class ScreenlogicEntity(CoordinatorEntity[ScreenlogicDataUpdateCoordinator]): """Base class for all ScreenLogic entities.""" - def __init__(self, coordinator, data_key, enabled=True): + def __init__( + self, + coordinator: ScreenlogicDataUpdateCoordinator, + data_key: str, + enabled: bool = True, + ) -> None: """Initialize of the entity.""" super().__init__(coordinator) self._data_key = data_key @@ -34,8 +40,10 @@ class ScreenlogicEntity(CoordinatorEntity[ScreenlogicDataUpdateCoordinator]): ] except KeyError: equipment_model = f"Unknown Model C:{controller_type} H:{hardware_type}" + mac = self.mac + assert mac is not None self._attr_device_info = DeviceInfo( - connections={(dr.CONNECTION_NETWORK_MAC, self.mac)}, + connections={(dr.CONNECTION_NETWORK_MAC, mac)}, manufacturer="Pentair", model=equipment_model, name=self.gateway_name, @@ -43,17 +51,18 @@ class ScreenlogicEntity(CoordinatorEntity[ScreenlogicDataUpdateCoordinator]): ) @property - def mac(self): + def mac(self) -> str | None: """Mac address.""" + assert self.coordinator.config_entry is not None return self.coordinator.config_entry.unique_id @property - def config_data(self): + def config_data(self) -> dict[str | int, Any]: """Shortcut for config data.""" return self.gateway_data[SL_DATA.KEY_CONFIG] @property - def gateway(self): + def gateway(self) -> ScreenLogicGateway: """Return the gateway.""" return self.coordinator.gateway @@ -63,18 +72,18 @@ class ScreenlogicEntity(CoordinatorEntity[ScreenlogicDataUpdateCoordinator]): return self.gateway.get_data() @property - def gateway_name(self): + def gateway_name(self) -> str: """Return the configured name of the gateway.""" return self.gateway.name - async def _async_refresh(self): + async def _async_refresh(self) -> None: """Refresh the data from the gateway.""" await self.coordinator.async_refresh() # Second debounced refresh to catch any secondary # changes in the device await self.coordinator.async_request_refresh() - async def _async_refresh_timed(self, now): + async def _async_refresh_timed(self, now: datetime) -> None: """Refresh from a timed called.""" await self.coordinator.async_request_refresh() @@ -82,7 +91,13 @@ class ScreenlogicEntity(CoordinatorEntity[ScreenlogicDataUpdateCoordinator]): class ScreenLogicPushEntity(ScreenlogicEntity): """Base class for all ScreenLogic push entities.""" - def __init__(self, coordinator, data_key, message_code, enabled=True): + def __init__( + self, + coordinator: ScreenlogicDataUpdateCoordinator, + data_key: str, + message_code: CODE, + enabled: bool = True, + ) -> None: """Initialize the entity.""" super().__init__(coordinator, data_key, enabled) self._update_message_code = message_code @@ -108,7 +123,7 @@ class ScreenLogicCircuitEntity(ScreenLogicPushEntity): _attr_has_entity_name = True @property - def name(self): + def name(self) -> str: """Get the name of the switch.""" return self.circuit["name"] @@ -117,15 +132,15 @@ class ScreenLogicCircuitEntity(ScreenLogicPushEntity): """Get whether the switch is in on state.""" return self.circuit["value"] == ON_OFF.ON - async def async_turn_on(self, **kwargs) -> None: + async def async_turn_on(self, **kwargs: Any) -> None: """Send the ON command.""" await self._async_set_circuit(ON_OFF.ON) - async def async_turn_off(self, **kwargs) -> None: + async def async_turn_off(self, **kwargs: Any) -> None: """Send the OFF command.""" await self._async_set_circuit(ON_OFF.OFF) - async def _async_set_circuit(self, circuit_value) -> None: + async def _async_set_circuit(self, circuit_value: int) -> None: if not await self.gateway.async_set_circuit(self._data_key, circuit_value): raise HomeAssistantError( f"Failed to set_circuit {self._data_key} {circuit_value}"