diff --git a/homeassistant/components/husqvarna_automower/coordinator.py b/homeassistant/components/husqvarna_automower/coordinator.py index 819ee41a43d..9456074596a 100644 --- a/homeassistant/components/husqvarna_automower/coordinator.py +++ b/homeassistant/components/husqvarna_automower/coordinator.py @@ -13,7 +13,7 @@ from aioautomower.exceptions import ( HusqvarnaTimeoutError, HusqvarnaWSServerHandshakeError, ) -from aioautomower.model import MowerAttributes +from aioautomower.model import MowerDictionary from aioautomower.session import AutomowerSession from homeassistant.config_entries import ConfigEntry @@ -32,7 +32,7 @@ DEFAULT_RECONNECT_TIME = 2 # Define a default reconnect time type AutomowerConfigEntry = ConfigEntry[AutomowerDataUpdateCoordinator] -class AutomowerDataUpdateCoordinator(DataUpdateCoordinator[dict[str, MowerAttributes]]): +class AutomowerDataUpdateCoordinator(DataUpdateCoordinator[MowerDictionary]): """Class to manage fetching Husqvarna data.""" config_entry: AutomowerConfigEntry @@ -61,7 +61,7 @@ class AutomowerDataUpdateCoordinator(DataUpdateCoordinator[dict[str, MowerAttrib self._zones_last_update: dict[str, set[str]] = {} self._areas_last_update: dict[str, set[int]] = {} - async def _async_update_data(self) -> dict[str, MowerAttributes]: + async def _async_update_data(self) -> MowerDictionary: """Subscribe for websocket and poll data from the API.""" if not self.ws_connected: await self.api.connect() @@ -84,7 +84,7 @@ class AutomowerDataUpdateCoordinator(DataUpdateCoordinator[dict[str, MowerAttrib return data @callback - def callback(self, ws_data: dict[str, MowerAttributes]) -> None: + def callback(self, ws_data: MowerDictionary) -> None: """Process websocket callbacks and write them to the DataUpdateCoordinator.""" self.async_set_updated_data(ws_data) @@ -119,7 +119,7 @@ class AutomowerDataUpdateCoordinator(DataUpdateCoordinator[dict[str, MowerAttrib "reconnect_task", ) - def _async_add_remove_devices(self, data: dict[str, MowerAttributes]) -> None: + def _async_add_remove_devices(self, data: MowerDictionary) -> None: """Add new device, remove non-existing device.""" current_devices = set(data) @@ -159,9 +159,7 @@ class AutomowerDataUpdateCoordinator(DataUpdateCoordinator[dict[str, MowerAttrib for mower_callback in self.new_devices_callbacks: mower_callback(new_devices) - def _async_add_remove_stay_out_zones( - self, data: dict[str, MowerAttributes] - ) -> None: + def _async_add_remove_stay_out_zones(self, data: MowerDictionary) -> None: """Add new stay-out zones, remove non-existing stay-out zones.""" current_zones = { mower_id: set(mower_data.stay_out_zones.zones) @@ -207,7 +205,7 @@ class AutomowerDataUpdateCoordinator(DataUpdateCoordinator[dict[str, MowerAttrib return current_zones - def _async_add_remove_work_areas(self, data: dict[str, MowerAttributes]) -> None: + def _async_add_remove_work_areas(self, data: MowerDictionary) -> None: """Add new work areas, remove non-existing work areas.""" current_areas = { mower_id: set(mower_data.work_areas)