Refactor Watts Vision+ to generic device, in preparation for switch support (#162721)

Co-authored-by: Joostlek <joostlek@outlook.com>

homeassistant/components/watts/__init__.py

@@ -20,9 +20,9 @@ from homeassistant.helpers.dispatcher import async_dispatcher_send
 
 from .const import DOMAIN
 from .coordinator import (
+    WattsVisionDeviceCoordinator,
+    WattsVisionDeviceData,
     WattsVisionHubCoordinator,
-    WattsVisionThermostatCoordinator,
-    WattsVisionThermostatData,
 )
 
 _LOGGER = logging.getLogger(__name__)
@@ -36,7 +36,7 @@ class WattsVisionRuntimeData:
 
     auth: WattsVisionAuth
     hub_coordinator: WattsVisionHubCoordinator
-    thermostat_coordinators: dict[str, WattsVisionThermostatCoordinator]
+    device_coordinators: dict[str, WattsVisionDeviceCoordinator]
     client: WattsVisionClient
 
 
@@ -44,46 +44,50 @@ type WattsVisionConfigEntry = ConfigEntry[WattsVisionRuntimeData]
 
 
 @callback
-def _handle_new_thermostats(
+def _handle_new_devices(
     hass: HomeAssistant,
     entry: WattsVisionConfigEntry,
     hub_coordinator: WattsVisionHubCoordinator,
 ) -> None:
-    """Check for new thermostat devices and create coordinators."""
+    """Check for new devices and create coordinators."""
-
     current_device_ids = set(hub_coordinator.data.keys())
-    known_device_ids = set(entry.runtime_data.thermostat_coordinators.keys())
+    known_device_ids = set(entry.runtime_data.device_coordinators.keys())
     new_device_ids = current_device_ids - known_device_ids
 
     if not new_device_ids:
         return
 
-    _LOGGER.info("Discovered %d new device(s): %s", len(new_device_ids), new_device_ids)
+    device_coordinators = entry.runtime_data.device_coordinators
-
-    thermostat_coordinators = entry.runtime_data.thermostat_coordinators
     client = entry.runtime_data.client
+    supported_device_ids: list[str] = []
 
     for device_id in new_device_ids:
         device = hub_coordinator.data[device_id]
         if not isinstance(device, ThermostatDevice):
             continue
 
-        thermostat_coordinator = WattsVisionThermostatCoordinator(
+        device_coordinator = WattsVisionDeviceCoordinator(
             hass, client, entry, hub_coordinator, device_id
         )
-        thermostat_coordinator.async_set_updated_data(
+        device_coordinator.async_set_updated_data(WattsVisionDeviceData(device=device))
-            WattsVisionThermostatData(thermostat=device)
+        device_coordinators[device_id] = device_coordinator
+        supported_device_ids.append(device_id)
+
+        _LOGGER.debug("Created device coordinator for device %s", device_id)
+
+    if not supported_device_ids:
+        return
+
+    _LOGGER.info(
+        "Discovered %d new device(s): %s",
+        len(supported_device_ids),
+        supported_device_ids,
     )
-        thermostat_coordinators[device_id] = thermostat_coordinator
-
-        _LOGGER.debug("Created thermostat coordinator for device %s", device_id)
-
     async_dispatcher_send(hass, f"{DOMAIN}_{entry.entry_id}_new_device")
 
 
 async def async_setup_entry(hass: HomeAssistant, entry: WattsVisionConfigEntry) -> bool:
     """Set up Watts Vision from a config entry."""
-
     try:
         implementation = (
             await config_entry_oauth2_flow.async_get_config_entry_implementation(
@@ -117,24 +121,22 @@ async def async_setup_entry(hass: HomeAssistant, entry: WattsVisionConfigEntry)
 
     await hub_coordinator.async_config_entry_first_refresh()
 
-    thermostat_coordinators: dict[str, WattsVisionThermostatCoordinator] = {}
+    device_coordinators: dict[str, WattsVisionDeviceCoordinator] = {}
     for device_id in hub_coordinator.device_ids:
         device = hub_coordinator.data[device_id]
         if not isinstance(device, ThermostatDevice):
             continue
 
-        thermostat_coordinator = WattsVisionThermostatCoordinator(
+        device_coordinator = WattsVisionDeviceCoordinator(
             hass, client, entry, hub_coordinator, device_id
         )
-        thermostat_coordinator.async_set_updated_data(
+        device_coordinator.async_set_updated_data(WattsVisionDeviceData(device=device))
-            WattsVisionThermostatData(thermostat=device)
+        device_coordinators[device_id] = device_coordinator
-        )
-        thermostat_coordinators[device_id] = thermostat_coordinator
 
     entry.runtime_data = WattsVisionRuntimeData(
         auth=auth,
         hub_coordinator=hub_coordinator,
-        thermostat_coordinators=thermostat_coordinators,
+        device_coordinators=device_coordinators,
         client=client,
     )
 
@@ -143,7 +145,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: WattsVisionConfigEntry)
     # Listener for dynamic device detection
     entry.async_on_unload(
         hub_coordinator.async_add_listener(
-            lambda: _handle_new_thermostats(hass, entry, hub_coordinator)
+            lambda: _handle_new_devices(hass, entry, hub_coordinator)
         )
     )
 
@@ -154,7 +156,7 @@ async def async_unload_entry(
     hass: HomeAssistant, entry: WattsVisionConfigEntry
 ) -> bool:
     """Unload a config entry."""
-    for thermostat_coordinator in entry.runtime_data.thermostat_coordinators.values():
+    for device_coordinator in entry.runtime_data.device_coordinators.values():
-        thermostat_coordinator.unsubscribe_hub_listener()
+        device_coordinator.unsubscribe_hub_listener()
 
     return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)

homeassistant/components/watts/climate.py

@@ -20,8 +20,8 @@ from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
 from . import WattsVisionConfigEntry
 from .const import DOMAIN, HVAC_MODE_TO_THERMOSTAT, THERMOSTAT_MODE_TO_HVAC
-from .coordinator import WattsVisionThermostatCoordinator
+from .coordinator import WattsVisionDeviceCoordinator
-from .entity import WattsVisionThermostatEntity
+from .entity import WattsVisionEntity
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -34,14 +34,18 @@ async def async_setup_entry(
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up Watts Vision climate entities from a config entry."""
-
+    device_coordinators = entry.runtime_data.device_coordinators
-    thermostat_coordinators = entry.runtime_data.thermostat_coordinators
     known_device_ids: set[str] = set()
 
     @callback
     def _check_new_thermostats() -> None:
         """Check for new thermostat devices."""
-        current_device_ids = set(thermostat_coordinators.keys())
+        thermostat_coords = {
+            did: coord
+            for did, coord in device_coordinators.items()
+            if isinstance(coord.data.device, ThermostatDevice)
+        }
+        current_device_ids = set(thermostat_coords.keys())
         new_device_ids = current_device_ids - known_device_ids
 
         if not new_device_ids:
@@ -52,13 +56,12 @@ async def async_setup_entry(
             len(new_device_ids),
         )
 
-        new_entities = [
+        new_entities = []
-            WattsVisionClimate(
+        for device_id in new_device_ids:
-                thermostat_coordinators[device_id],
+            coord = thermostat_coords[device_id]
-                thermostat_coordinators[device_id].data.thermostat,
+            device = coord.data.device
-            )
+            assert isinstance(device, ThermostatDevice)
-            for device_id in new_device_ids
+            new_entities.append(WattsVisionClimate(coord, device))
-        ]
 
         known_device_ids.update(new_device_ids)
         async_add_entities(new_entities)
@@ -75,7 +78,7 @@ async def async_setup_entry(
     )
 
 
-class WattsVisionClimate(WattsVisionThermostatEntity, ClimateEntity):
+class WattsVisionClimate(WattsVisionEntity[ThermostatDevice], ClimateEntity):
     """Representation of a Watts Vision heater as a climate entity."""
 
     _attr_supported_features = ClimateEntityFeature.TARGET_TEMPERATURE
@@ -84,11 +87,10 @@ class WattsVisionClimate(WattsVisionThermostatEntity, ClimateEntity):
 
     def __init__(
         self,
-        coordinator: WattsVisionThermostatCoordinator,
+        coordinator: WattsVisionDeviceCoordinator,
         thermostat: ThermostatDevice,
     ) -> None:
         """Initialize the climate entity."""
-
         super().__init__(coordinator, thermostat.device_id)
 
         self._attr_min_temp = thermostat.min_allowed_temperature
@@ -102,17 +104,17 @@ class WattsVisionClimate(WattsVisionThermostatEntity, ClimateEntity):
     @property
     def current_temperature(self) -> float | None:
         """Return the current temperature."""
-        return self.thermostat.current_temperature
+        return self.device.current_temperature
 
     @property
     def target_temperature(self) -> float | None:
         """Return the temperature setpoint."""
-        return self.thermostat.setpoint
+        return self.device.setpoint
 
     @property
     def hvac_mode(self) -> HVACMode | None:
         """Return hvac mode."""
-        return THERMOSTAT_MODE_TO_HVAC.get(self.thermostat.thermostat_mode)
+        return THERMOSTAT_MODE_TO_HVAC.get(self.device.thermostat_mode)
 
     async def async_set_temperature(self, **kwargs: Any) -> None:
         """Set new target temperature."""

homeassistant/components/watts/coordinator.py

@@ -15,7 +15,7 @@ from visionpluspython.exceptions import (
     WattsVisionError,
     WattsVisionTimeoutError,
 )
-from visionpluspython.models import Device, ThermostatDevice
+from visionpluspython.models import Device
 
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
@@ -39,10 +39,10 @@ _LOGGER = logging.getLogger(__name__)
 
 
 @dataclass
-class WattsVisionThermostatData:
+class WattsVisionDeviceData:
-    """Data class for thermostat device coordinator."""
+    """Data class for device coordinator."""
 
-    thermostat: ThermostatDevice
+    device: Device
 
 
 class WattsVisionHubCoordinator(DataUpdateCoordinator[dict[str, Device]]):
@@ -150,10 +150,8 @@ class WattsVisionHubCoordinator(DataUpdateCoordinator[dict[str, Device]]):
         return list((self.data or {}).keys())
 
 
-class WattsVisionThermostatCoordinator(
+class WattsVisionDeviceCoordinator(DataUpdateCoordinator[WattsVisionDeviceData]):
-    DataUpdateCoordinator[WattsVisionThermostatData]
+    """Device coordinator for individual updates."""
-):
-    """Thermostat device coordinator for individual updates."""
 
     def __init__(
         self,
@@ -163,7 +161,7 @@ class WattsVisionThermostatCoordinator(
         hub_coordinator: WattsVisionHubCoordinator,
         device_id: str,
     ) -> None:
-        """Initialize the thermostat coordinator."""
+        """Initialize the device coordinator."""
         super().__init__(
             hass,
             _LOGGER,
@@ -185,11 +183,10 @@ class WattsVisionThermostatCoordinator(
         """Handle updates from hub coordinator."""
         if self.hub_coordinator.data and self.device_id in self.hub_coordinator.data:
             device = self.hub_coordinator.data[self.device_id]
-            assert isinstance(device, ThermostatDevice)
+            self.async_set_updated_data(WattsVisionDeviceData(device=device))
-            self.async_set_updated_data(WattsVisionThermostatData(thermostat=device))
 
-    async def _async_update_data(self) -> WattsVisionThermostatData:
+    async def _async_update_data(self) -> WattsVisionDeviceData:
-        """Refresh specific thermostat device."""
+        """Refresh specific device."""
         if self._fast_polling_until and datetime.now() > self._fast_polling_until:
             self._fast_polling_until = None
             self.update_interval = None
@@ -215,9 +212,8 @@ class WattsVisionThermostatCoordinator(
         if not device:
             raise UpdateFailed(f"Device {self.device_id} not found")
 
-        assert isinstance(device, ThermostatDevice)
+        _LOGGER.debug("Refreshed device %s", self.device_id)
-        _LOGGER.debug("Refreshed thermostat %s", self.device_id)
+        return WattsVisionDeviceData(device=device)
-        return WattsVisionThermostatData(thermostat=device)
 
     def trigger_fast_polling(self, duration: int = 60) -> None:
         """Activate fast polling for a specified duration after a command."""

homeassistant/components/watts/entity.py

@@ -2,42 +2,44 @@
 
 from __future__ import annotations
 
-from visionpluspython.models import ThermostatDevice
+from typing import cast
+
+from visionpluspython.models import Device
 
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
 from .const import DOMAIN
-from .coordinator import WattsVisionThermostatCoordinator
+from .coordinator import WattsVisionDeviceCoordinator
 
 
-class WattsVisionThermostatEntity(CoordinatorEntity[WattsVisionThermostatCoordinator]):
+class WattsVisionEntity[_T: Device](CoordinatorEntity[WattsVisionDeviceCoordinator]):
-    """Base entity for Watts Vision thermostat devices."""
+    """Base entity for Watts Vision devices."""
 
     _attr_has_entity_name = True
 
     def __init__(
-        self, coordinator: WattsVisionThermostatCoordinator, device_id: str
+        self, coordinator: WattsVisionDeviceCoordinator, device_id: str
     ) -> None:
         """Initialize the entity."""
-
         super().__init__(coordinator, context=device_id)
         self.device_id = device_id
         self._attr_unique_id = device_id
+        device = coordinator.data.device
         self._attr_device_info = DeviceInfo(
             identifiers={(DOMAIN, self.device_id)},
-            name=self.thermostat.device_name,
+            name=device.device_name,
             manufacturer="Watts",
-            model=f"Vision+ {self.thermostat.device_type}",
+            model=f"Vision+ {device.device_type}",
-            suggested_area=self.thermostat.room_name,
+            suggested_area=device.room_name,
         )
 
     @property
-    def thermostat(self) -> ThermostatDevice:
+    def device(self) -> _T:
-        """Return the thermostat device from the coordinator data."""
+        """Return the device from the coordinator data."""
-        return self.coordinator.data.thermostat
+        return cast(_T, self.coordinator.data.device)
 
     @property
     def available(self) -> bool:
         """Return True if entity is available."""
-        return super().available and self.coordinator.data.thermostat.is_online
+        return super().available and self.coordinator.data.device.is_online