Adapt relevant code

Author: bouwew

custom_components/plugwise/binary_sensor.py

@@ -174,7 +174,7 @@ def is_on(self) -> bool:
                     self.hass, message, "Plugwise Notification:", f"{DOMAIN}.{notify_id}"
                 )
 
-        return self.device[BINARY_SENSORS][self.entity_description.key]
+        return self.device_or_zone[BINARY_SENSORS][self.entity_description.key]
 
     @property
     def extra_state_attributes(self) -> Mapping[str, Any] | None:

custom_components/plugwise/climate.py

@@ -76,19 +76,30 @@ async def async_setup_entry(
     @callback
     def _add_entities() -> None:
         """Add Entities during init and runtime."""
-        if not coordinator.new_devices:
+        if not (coordinator.new_devices or coordinator.new_zones):
             return
 
         entities: list[PlugwiseClimateEntity] = []
-        for device_id in coordinator.new_devices:
-            device = coordinator.data.devices[device_id]
-            if device[DEV_CLASS] in MASTER_THERMOSTATS:
+        if coordinator.new_zones:
+            for device_id in coordinator.new_zones:
+                thermostat = coordinator.data.zones[device_id]
                 entities.append(
                     PlugwiseClimateEntity(
                         coordinator, device_id, homekit_enabled
                     )  # pw-beta homekit emulation
                 )
-                LOGGER.debug("Add climate %s", device[ATTR_NAME])
+                LOGGER.debug("Add climate %s", thermostat[ATTR_NAME])
+
+        if coordinator.new_devices:
+            for device_id in coordinator.new_devices:
+                device = coordinator.data.devices[device_id]
+                if device[DEV_CLASS] in MASTER_THERMOSTATS:
+                    entities.append(
+                        PlugwiseClimateEntity(
+                            coordinator, device_id, homekit_enabled
+                        )  # pw-beta homekit emulation
+                    )
+                    LOGGER.debug("Add climate %s", device[ATTR_NAME])
 
         async_add_entities(entities)
 
@@ -118,14 +129,18 @@ def __init__(
         super().__init__(coordinator, device_id)
 
         self._homekit_enabled = homekit_enabled  # pw-beta homekit emulation
+        self._location = device_id
+        if (location := self.device_or_zone.get(LOCATION)) is not None:
+            self._location = location
+
         gateway_id: str = coordinator.data.gateway[GATEWAY_ID]
         self.gateway_data = coordinator.data.devices[gateway_id]
 
-        self._attr_max_temp = min(self.device[THERMOSTAT][UPPER_BOUND], 35.0)
-        self._attr_min_temp = self.device[THERMOSTAT][LOWER_BOUND]
+        self._attr_max_temp = min(self.device_or_zone[THERMOSTAT][UPPER_BOUND], 35.0)
+        self._attr_min_temp = self.device_or_zone[THERMOSTAT][LOWER_BOUND]
         # Ensure we don't drop below 0.1
         self._attr_target_temperature_step = max(
-            self.device[THERMOSTAT][RESOLUTION], 0.1
+            self.device_or_zone[THERMOSTAT][RESOLUTION], 0.1
         )
         self._attr_unique_id = f"{device_id}-climate"
 
@@ -143,7 +158,7 @@ def __init__(
             self._attr_supported_features |= (
                 ClimateEntityFeature.TURN_OFF | ClimateEntityFeature.TURN_ON
             )
-        if presets := self.device["preset_modes"]:  # can be NONE
+        if presets := self.device_or_zone["preset_modes"]:  # can be NONE
             self._attr_supported_features |= ClimateEntityFeature.PRESET_MODE
         self._attr_preset_modes = presets
 
@@ -164,7 +179,7 @@ def _previous_action_mode(self, coordinator: PlugwiseDataUpdateCoordinator) -> N
     @property
     def current_temperature(self) -> float:
         """Return the current temperature."""
-        return self.device[SENSORS][ATTR_TEMPERATURE]
+        return self.device_or_zone[SENSORS][ATTR_TEMPERATURE]
 
     @property
     def target_temperature(self) -> float:
@@ -173,29 +188,29 @@ def target_temperature(self) -> float:
         Connected to the HVACMode combination of AUTO-HEAT.
         """
 
-        return self.device[THERMOSTAT][TARGET_TEMP]
+        return self.device_or_zone[THERMOSTAT][TARGET_TEMP]
 
     @property
     def target_temperature_high(self) -> float:
         """Return the temperature we try to reach in case of cooling.
 
         Connected to the HVACMode combination of AUTO-HEAT_COOL.
         """
-        return self.device[THERMOSTAT][TARGET_TEMP_HIGH]
+        return self.device_or_zone[THERMOSTAT][TARGET_TEMP_HIGH]
 
     @property
     def target_temperature_low(self) -> float:
         """Return the heating temperature we try to reach in case of heating.
 
         Connected to the HVACMode combination AUTO-HEAT_COOL.
         """
-        return self.device[THERMOSTAT][TARGET_TEMP_LOW]
+        return self.device_or_zone[THERMOSTAT][TARGET_TEMP_LOW]
 
     @property
     def hvac_mode(self) -> HVACMode:
         """Return HVAC operation ie. auto, cool, heat, heat_cool, or off mode."""
         if (
-            mode := self.device[CLIMATE_MODE]
+            mode := self.device_or_zone[CLIMATE_MODE]
         ) is None or mode not in self.hvac_modes:  # pw-beta add to Core
             return HVACMode.HEAT  # pragma: no cover
         # pw-beta homekit emulation
@@ -214,7 +229,7 @@ def hvac_modes(self) -> list[HVACMode]:
         ):
             hvac_modes.append(HVACMode.OFF)
 
-        if AVAILABLE_SCHEDULES in self.device:
+        if AVAILABLE_SCHEDULES in self.device_or_zone:
             hvac_modes.append(HVACMode.AUTO)
 
         if self.cdr_gateway[COOLING_PRESENT]:
@@ -237,7 +252,7 @@ def hvac_action(self) -> HVACAction:  # pw-beta add to Core
         self._previous_action_mode(self.coordinator)
 
         # Adam provides the hvac_action for each thermostat
-        if (control_state := self.device.get(CONTROL_STATE)) in (HVACAction.COOLING, HVACAction.HEATING, HVACAction.PREHEATING):
+        if (control_state := self.device_or_zone.get(CONTROL_STATE)) in (HVACAction.COOLING, HVACAction.HEATING, HVACAction.PREHEATING):
             return cast(HVACAction, control_state)
         if control_state == HVACMode.OFF:
             return HVACAction.IDLE
@@ -255,7 +270,7 @@ def hvac_action(self) -> HVACAction:  # pw-beta add to Core
     @property
     def preset_mode(self) -> str | None:
         """Return the current preset mode."""
-        return self.device[ACTIVE_PRESET]
+        return self.device_or_zone[ACTIVE_PRESET]
 
     @plugwise_command
     async def async_set_temperature(self, **kwargs: Any) -> None:
@@ -273,7 +288,7 @@ async def async_set_temperature(self, **kwargs: Any) -> None:
         if mode := kwargs.get(ATTR_HVAC_MODE):
             await self.async_set_hvac_mode(mode)
 
-        await self.coordinator.api.set_temperature(self.device[LOCATION], data)
+        await self.coordinator.api.set_temperature(self._location, data)
 
     @plugwise_command
     async def async_set_hvac_mode(self, hvac_mode: HVACMode) -> None:
@@ -286,7 +301,7 @@ async def async_set_hvac_mode(self, hvac_mode: HVACMode) -> None:
 
         if hvac_mode != HVACMode.OFF:
             await self.coordinator.api.set_schedule_state(
-                self.device[LOCATION],
+                self._location,
                 STATE_ON if hvac_mode == HVACMode.AUTO else STATE_OFF,
             )
 
@@ -303,11 +318,11 @@ async def async_set_hvac_mode(self, hvac_mode: HVACMode) -> None:
                 await self.async_set_preset_mode(PRESET_AWAY)  # pragma: no cover
             if (
                 self._homekit_mode in [HVACMode.HEAT, HVACMode.HEAT_COOL]
-                and self.device[ACTIVE_PRESET] == PRESET_AWAY
+                and self.device_or_zone[ACTIVE_PRESET] == PRESET_AWAY
             ):  # pragma: no cover
                 await self.async_set_preset_mode(PRESET_HOME)  # pragma: no cover
 
     @plugwise_command
     async def async_set_preset_mode(self, preset_mode: str) -> None:
         """Set the preset mode."""
-        await self.coordinator.api.set_preset(self.device[LOCATION], preset_mode)
+        await self.coordinator.api.set_preset(self._location, preset_mode)

custom_components/plugwise/coordinator.py

@@ -70,7 +70,9 @@ def __init__(
             websession=async_get_clientsession(hass, verify_ssl=False),
         )
         self._current_devices: set[str] = set()
+        self._current_zones: set[str] = set()
         self.new_devices: set[str] = set()
+        self.new_zones: set[str] = set()
         self.update_interval = update_interval
 
     async def _connect(self) -> None:
@@ -91,7 +93,7 @@ async def _connect(self) -> None:
 
     async def _async_update_data(self) -> PlugwiseData:
         """Fetch data from Plugwise."""
-        data = PlugwiseData({}, {})
+        data = PlugwiseData(devices={}, gateway={}, zones={})
         try:
             if not self._connected:
                 await self._connect()
@@ -110,19 +112,24 @@ async def _async_update_data(self) -> PlugwiseData:
             raise ConfigEntryError("Device with unsupported firmware") from err
         else:
             LOGGER.debug(f"{self.api.smile_name} data: %s", data)
-            await self.async_add_remove_devices(data, self.config_entry)
+            await self.async_add_remove_devices_zones(data, self.config_entry)
 
         return data
 
-    async def async_add_remove_devices(self, data: PlugwiseData, entry: ConfigEntry) -> None:
+    async def async_add_remove_devices_zones(self, data: PlugwiseData, entry: ConfigEntry) -> None:
         """Add new Plugwise devices, remove non-existing devices."""
         # Check for new or removed devices
         self.new_devices = set(data.devices) - self._current_devices
+        self.new_zones = set(data.zones) - self._current_zones
         removed_devices = self._current_devices - set(data.devices)
+        removed_zones = self._current_zones - set(data.zones)
         self._current_devices = set(data.devices)
+        self._current_zones = set(data.zones)
 
         if removed_devices:
             await self.async_remove_devices(data, entry)
+        # if removed_zones:
+            # await self.async_remove_zones(data, entry)
 
     async def async_remove_devices(self, data: PlugwiseData, entry: ConfigEntry) -> None:
         """Clean registries when removed devices found."""

custom_components/plugwise/diagnostics.py

@@ -17,4 +17,5 @@ async def async_get_config_entry_diagnostics(
     return {
         "gateway": coordinator.data.gateway,
         "devices": coordinator.data.devices,
+        "zones": coordinator.data.zones,
     }

custom_components/plugwise/entity.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from plugwise.constants import DeviceData
+from plugwise.constants import DeviceZoneData
 
 from homeassistant.const import ATTR_NAME, ATTR_VIA_DEVICE, CONF_HOST
 from homeassistant.helpers.device_registry import (
@@ -48,51 +48,54 @@ def __init__(
         if entry := self.coordinator.config_entry:
             configuration_url = f"http://{entry.data[CONF_HOST]}"
 
-        data = coordinator.data.devices[device_id]
-        connections = set()
-        if mac := data.get(MAC_ADDRESS):
-            connections.add((CONNECTION_NETWORK_MAC, mac))
-        if mac := data.get(ZIGBEE_MAC_ADDRESS):
-            connections.add((CONNECTION_ZIGBEE, mac))
-
-        self._attr_device_info = DeviceInfo(
-            configuration_url=configuration_url,
-            identifiers={(DOMAIN, device_id)},
-            connections=connections,
-            manufacturer=data.get(VENDOR),
-            model=data.get(MODEL),
-            model_id=data.get(MODEL_ID),
-            name=coordinator.data.gateway[SMILE_NAME],
-            sw_version=data.get(FIRMWARE),
-            hw_version=data.get(HARDWARE),
-        )
-
-        if device_id != coordinator.data.gateway[GATEWAY_ID]:
-            self._attr_device_info.update(
-                {
-                    ATTR_NAME: data.get(ATTR_NAME),
-                    ATTR_VIA_DEVICE: (
-                        DOMAIN,
-                        str(self.coordinator.data.gateway[GATEWAY_ID]),
-                    ),
-                }
+        if (data := coordinator.data.devices.get(device_id)) is not None:
+            connections = set()
+            if mac := data.get(MAC_ADDRESS):
+                connections.add((CONNECTION_NETWORK_MAC, mac))
+            if mac := data.get(ZIGBEE_MAC_ADDRESS):
+                connections.add((CONNECTION_ZIGBEE, mac))
+
+            self._attr_device_info = DeviceInfo(
+                configuration_url=configuration_url,
+                identifiers={(DOMAIN, device_id)},
+                connections=connections,
+                manufacturer=data.get(VENDOR),
+                model=data.get(MODEL),
+                model_id=data.get(MODEL_ID),
+                name=coordinator.data.gateway[SMILE_NAME],
+                sw_version=data.get(FIRMWARE),
+                hw_version=data.get(HARDWARE),
             )
 
+            if device_id != coordinator.data.gateway[GATEWAY_ID]:
+                self._attr_device_info.update(
+                    {
+                        ATTR_NAME: data.get(ATTR_NAME),
+                        ATTR_VIA_DEVICE: (
+                            DOMAIN,
+                            str(self.coordinator.data.gateway[GATEWAY_ID]),
+                        ),
+                    }
+                )
+
     @property
     def available(self) -> bool:
         """Return if entity is available."""
         return (
             # Upstream: Do not change the AVAILABLE line below: some Plugwise devices
             # Upstream: do not provide their availability-status!
             self._dev_id in self.coordinator.data.devices
-            and (AVAILABLE not in self.device or self.device[AVAILABLE] is True)
+            and (AVAILABLE not in self.device_or_zone or self.device_or_zone[AVAILABLE] is True)
             and super().available
         )
 
     @property
-    def device(self) -> DeviceData:
-        """Return data for this device."""
-        return self.coordinator.data.devices[self._dev_id]
+    def device_or_zone(self) -> DeviceZoneData:
+        """Return the device or zone connected to the dev_id."""
+        if (device := self.coordinator.data.devices[self._dev_id]) is not None:
+            return device
+
+        return self.coordinator.data.zones[self._dev_id]
 
     async def async_added_to_hass(self) -> None:
         """Subscribe to updates."""

custom_components/plugwise/number.py

@@ -122,23 +122,23 @@ def __init__(
     ) -> None:
         """Initiate Plugwise Number."""
         super().__init__(coordinator, device_id)
-        self.actuator = self.device[description.key]  # Upstream
+        self.actuator = self.device_or_zone[description.key]  # Upstream
         self.device_id = device_id
         self.entity_description = description
         self._attr_unique_id = f"{device_id}-{description.key}"
         self._attr_mode = NumberMode.BOX
-        self._attr_native_max_value = self.device[description.key][UPPER_BOUND]  # Upstream const
-        self._attr_native_min_value = self.device[description.key][LOWER_BOUND]  # Upstream const
+        self._attr_native_max_value = self.device_or_zone[description.key][UPPER_BOUND]  # Upstream const
+        self._attr_native_min_value = self.device_or_zone[description.key][LOWER_BOUND]  # Upstream const
 
-        native_step = self.device[description.key][RESOLUTION]  # Upstream const
+        native_step = self.device_or_zone[description.key][RESOLUTION]  # Upstream const
         if description.key != TEMPERATURE_OFFSET:  # Upstream const
             native_step = max(native_step, 0.5)
         self._attr_native_step = native_step
 
     @property
     def native_value(self) -> float:
         """Return the present setpoint value."""
-        return self.device[self.entity_description.key]["setpoint"]
+        return self.device_or_zone[self.entity_description.key]["setpoint"]
 
     @plugwise_command
     async def async_set_native_value(self, value: float) -> None:

custom_components/plugwise/select.py

@@ -124,27 +124,30 @@ def __init__(
     ) -> None:
         """Initialise the selector."""
         super().__init__(coordinator, device_id)
-        self.entity_description = entity_description
         self._attr_unique_id = f"{device_id}-{entity_description.key}"
+        self.entity_description = entity_description
+        self._location = device_id
+        if (location := self.device_or_zone.get(LOCATION)) is not None:
+            self._location = location
 
     @property
     def current_option(self) -> str:
         """Return the selected entity option to represent the entity state."""
-        return self.device[self.entity_description.key]
+        return self.device_or_zone[self.entity_description.key]
 
     @property
     def options(self) -> list[str]:
         """Return the available select-options."""
-        return self.device[self.entity_description.options_key]
+        return self.device_or_zone[self.entity_description.options_key]
 
     @plugwise_command
     async def async_select_option(self, option: str) -> None:
         """Change to the selected entity option.
 
-        self.device[LOCATION] and STATE_ON are required for the thermostat-schedule select.
+        Location ID and STATE_ON are required for the thermostat-schedule select.
         """
         await self.coordinator.api.set_select(
-            self.entity_description.key, self.device[LOCATION], option, STATE_ON
+            self.entity_description.key, self._location, option, STATE_ON
         )
         LOGGER.debug(
             "Set %s to %s was successful",