diff --git a/homeassistant/components/rainmachine/__init__.py b/homeassistant/components/rainmachine/__init__.py index 3feeac7a827..c30ce81dc6d 100644 --- a/homeassistant/components/rainmachine/__init__.py +++ b/homeassistant/components/rainmachine/__init__.py @@ -2,9 +2,10 @@ from __future__ import annotations import asyncio +from dataclasses import dataclass from datetime import timedelta from functools import partial -from typing import Any, cast +from typing import Any from regenmaschine import Client from regenmaschine.controller import Controller @@ -28,7 +29,7 @@ from homeassistant.helpers import ( device_registry as dr, entity_registry as er, ) -from homeassistant.helpers.entity import DeviceInfo, EntityDescription +from homeassistant.helpers.entity import DeviceInfo from homeassistant.helpers.update_coordinator import ( CoordinatorEntity, DataUpdateCoordinator, @@ -39,8 +40,6 @@ from homeassistant.util.network import is_ip_address from .config_flow import get_client_controller from .const import ( CONF_ZONE_RUN_TIME, - DATA_CONTROLLER, - DATA_COORDINATOR, DATA_PROGRAMS, DATA_PROVISION_SETTINGS, DATA_RESTRICTIONS_CURRENT, @@ -49,6 +48,7 @@ from .const import ( DOMAIN, LOGGER, ) +from .model import RainMachineEntityDescription DEFAULT_SSL = True @@ -135,6 +135,14 @@ SERVICE_RESTRICT_WATERING_SCHEMA = SERVICE_SCHEMA.extend( ) +@dataclass +class RainMachineData: + """Define an object to be stored in `hass.data`.""" + + controller: Controller + coordinators: dict[str, DataUpdateCoordinator] + + @callback def async_get_controller_for_service_call( hass: HomeAssistant, call: ServiceCall @@ -146,9 +154,8 @@ def async_get_controller_for_service_call( if device_entry := device_registry.async_get(device_id): for entry in hass.config_entries.async_entries(DOMAIN): if entry.entry_id in device_entry.config_entries: - return cast( - Controller, hass.data[DOMAIN][entry.entry_id][DATA_CONTROLLER] - ) + data: RainMachineData = hass.data[DOMAIN][entry.entry_id] + return data.controller raise ValueError(f"No controller for device ID: {device_id}") @@ -161,14 +168,12 @@ async def async_update_programs_and_zones( Program and zone updates always go together because of how linked they are: programs affect zones and certain combinations of zones affect programs. """ + data: RainMachineData = hass.data[DOMAIN][entry.entry_id] + await asyncio.gather( *[ - hass.data[DOMAIN][entry.entry_id][DATA_COORDINATOR][ - DATA_PROGRAMS - ].async_refresh(), - hass.data[DOMAIN][entry.entry_id][DATA_COORDINATOR][ - DATA_ZONES - ].async_refresh(), + data.coordinators[DATA_PROGRAMS].async_refresh(), + data.coordinators[DATA_ZONES].async_refresh(), ] ) @@ -250,10 +255,9 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: await asyncio.gather(*controller_init_tasks) hass.data.setdefault(DOMAIN, {}) - hass.data[DOMAIN][entry.entry_id] = { - DATA_CONTROLLER: controller, - DATA_COORDINATOR: coordinators, - } + hass.data[DOMAIN][entry.entry_id] = RainMachineData( + controller=controller, coordinators=coordinators + ) await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS) @@ -406,28 +410,27 @@ class RainMachineEntity(CoordinatorEntity): def __init__( self, entry: ConfigEntry, - coordinator: DataUpdateCoordinator, - controller: Controller, - description: EntityDescription, + data: RainMachineData, + description: RainMachineEntityDescription, ) -> None: """Initialize.""" - super().__init__(coordinator) + super().__init__(data.coordinators[description.api_category]) self._attr_device_info = DeviceInfo( - identifiers={(DOMAIN, controller.mac)}, + identifiers={(DOMAIN, data.controller.mac)}, configuration_url=f"https://{entry.data[CONF_IP_ADDRESS]}:{entry.data[CONF_PORT]}", - connections={(dr.CONNECTION_NETWORK_MAC, controller.mac)}, - name=str(controller.name).capitalize(), + connections={(dr.CONNECTION_NETWORK_MAC, data.controller.mac)}, + name=str(data.controller.name).capitalize(), manufacturer="RainMachine", model=( - f"Version {controller.hardware_version} " - f"(API: {controller.api_version})" + f"Version {data.controller.hardware_version} " + f"(API: {data.controller.api_version})" ), - sw_version=controller.software_version, + sw_version=data.controller.software_version, ) self._attr_extra_state_attributes = {} - self._attr_unique_id = f"{controller.mac}_{description.key}" - self._controller = controller + self._attr_unique_id = f"{data.controller.mac}_{description.key}" + self._data = data self.entity_description = description @callback diff --git a/homeassistant/components/rainmachine/binary_sensor.py b/homeassistant/components/rainmachine/binary_sensor.py index 6ba374a28ba..d9448d68f9d 100644 --- a/homeassistant/components/rainmachine/binary_sensor.py +++ b/homeassistant/components/rainmachine/binary_sensor.py @@ -10,16 +10,17 @@ from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.entity import EntityCategory from homeassistant.helpers.entity_platform import AddEntitiesCallback -from . import RainMachineEntity +from . import RainMachineData, RainMachineEntity from .const import ( - DATA_CONTROLLER, - DATA_COORDINATOR, DATA_PROVISION_SETTINGS, DATA_RESTRICTIONS_CURRENT, DATA_RESTRICTIONS_UNIVERSAL, DOMAIN, ) -from .model import RainMachineDescriptionMixinApiCategory +from .model import ( + RainMachineEntityDescription, + RainMachineEntityDescriptionMixinDataKey, +) from .util import key_exists TYPE_FLOW_SENSOR = "flow_sensor" @@ -35,7 +36,9 @@ TYPE_WEEKDAY = "weekday" @dataclass class RainMachineBinarySensorDescription( - BinarySensorEntityDescription, RainMachineDescriptionMixinApiCategory + BinarySensorEntityDescription, + RainMachineEntityDescription, + RainMachineEntityDescriptionMixinDataKey, ): """Describe a RainMachine binary sensor.""" @@ -124,8 +127,7 @@ async def async_setup_entry( hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback ) -> None: """Set up RainMachine binary sensors based on a config entry.""" - controller = hass.data[DOMAIN][entry.entry_id][DATA_CONTROLLER] - coordinators = hass.data[DOMAIN][entry.entry_id][DATA_COORDINATOR] + data: RainMachineData = hass.data[DOMAIN][entry.entry_id] api_category_sensor_map = { DATA_PROVISION_SETTINGS: ProvisionSettingsBinarySensor, @@ -135,12 +137,10 @@ async def async_setup_entry( async_add_entities( [ - api_category_sensor_map[description.api_category]( - entry, coordinator, controller, description - ) + api_category_sensor_map[description.api_category](entry, data, description) for description in BINARY_SENSOR_DESCRIPTIONS if ( - (coordinator := coordinators[description.api_category]) is not None + (coordinator := data.coordinators[description.api_category]) is not None and coordinator.data and key_exists(coordinator.data, description.data_key) ) @@ -151,6 +151,8 @@ async def async_setup_entry( class CurrentRestrictionsBinarySensor(RainMachineEntity, BinarySensorEntity): """Define a binary sensor that handles current restrictions data.""" + entity_description: RainMachineBinarySensorDescription + @callback def update_from_latest_data(self) -> None: """Update the state.""" @@ -171,6 +173,8 @@ class CurrentRestrictionsBinarySensor(RainMachineEntity, BinarySensorEntity): class ProvisionSettingsBinarySensor(RainMachineEntity, BinarySensorEntity): """Define a binary sensor that handles provisioning data.""" + entity_description: RainMachineBinarySensorDescription + @callback def update_from_latest_data(self) -> None: """Update the state.""" @@ -181,6 +185,8 @@ class ProvisionSettingsBinarySensor(RainMachineEntity, BinarySensorEntity): class UniversalRestrictionsBinarySensor(RainMachineEntity, BinarySensorEntity): """Define a binary sensor that handles universal restrictions data.""" + entity_description: RainMachineBinarySensorDescription + @callback def update_from_latest_data(self) -> None: """Update the state.""" diff --git a/homeassistant/components/rainmachine/const.py b/homeassistant/components/rainmachine/const.py index 56c1660a0ba..f94e7011dce 100644 --- a/homeassistant/components/rainmachine/const.py +++ b/homeassistant/components/rainmachine/const.py @@ -7,8 +7,6 @@ DOMAIN = "rainmachine" CONF_ZONE_RUN_TIME = "zone_run_time" -DATA_CONTROLLER = "controller" -DATA_COORDINATOR = "coordinator" DATA_PROGRAMS = "programs" DATA_PROVISION_SETTINGS = "provision.settings" DATA_RESTRICTIONS_CURRENT = "restrictions.current" diff --git a/homeassistant/components/rainmachine/diagnostics.py b/homeassistant/components/rainmachine/diagnostics.py index e5e249d41f4..58b918c18ee 100644 --- a/homeassistant/components/rainmachine/diagnostics.py +++ b/homeassistant/components/rainmachine/diagnostics.py @@ -3,15 +3,13 @@ from __future__ import annotations from typing import Any -from regenmaschine.controller import Controller - from homeassistant.components.diagnostics import async_redact_data from homeassistant.config_entries import ConfigEntry from homeassistant.const import CONF_LATITUDE, CONF_LONGITUDE, CONF_PASSWORD from homeassistant.core import HomeAssistant -from homeassistant.helpers.update_coordinator import DataUpdateCoordinator -from .const import DATA_CONTROLLER, DATA_COORDINATOR, DOMAIN +from . import RainMachineData +from .const import DOMAIN TO_REDACT = { CONF_LATITUDE, @@ -24,9 +22,7 @@ async def async_get_config_entry_diagnostics( hass: HomeAssistant, entry: ConfigEntry ) -> dict[str, Any]: """Return diagnostics for a config entry.""" - data = hass.data[DOMAIN][entry.entry_id] - coordinators: dict[str, DataUpdateCoordinator] = data[DATA_COORDINATOR] - controller: Controller = data[DATA_CONTROLLER] + data: RainMachineData = hass.data[DOMAIN][entry.entry_id] return { "entry": { @@ -38,15 +34,15 @@ async def async_get_config_entry_diagnostics( "coordinator": async_redact_data( { api_category: controller.data - for api_category, controller in coordinators.items() + for api_category, controller in data.coordinators.items() }, TO_REDACT, ), "controller": { - "api_version": controller.api_version, - "hardware_version": controller.hardware_version, - "name": controller.name, - "software_version": controller.software_version, + "api_version": data.controller.api_version, + "hardware_version": data.controller.hardware_version, + "name": data.controller.name, + "software_version": data.controller.software_version, }, }, } diff --git a/homeassistant/components/rainmachine/model.py b/homeassistant/components/rainmachine/model.py index 680a47c5d42..9ae99fe247a 100644 --- a/homeassistant/components/rainmachine/model.py +++ b/homeassistant/components/rainmachine/model.py @@ -1,17 +1,32 @@ """Define RainMachine data models.""" from dataclasses import dataclass +from homeassistant.helpers.entity import EntityDescription + @dataclass -class RainMachineDescriptionMixinApiCategory: - """Define an entity description mixin for binary and regular sensors.""" +class RainMachineEntityDescriptionMixinApiCategory: + """Define an entity description mixin to include an API category.""" api_category: str + + +@dataclass +class RainMachineEntityDescriptionMixinDataKey: + """Define an entity description mixin to include a data payload key.""" + data_key: str @dataclass -class RainMachineDescriptionMixinUid: - """Define an entity description mixin for switches.""" +class RainMachineEntityDescriptionMixinUid: + """Define an entity description mixin to include an activity UID.""" uid: int + + +@dataclass +class RainMachineEntityDescription( + EntityDescription, RainMachineEntityDescriptionMixinApiCategory +): + """Describe a RainMachine entity.""" diff --git a/homeassistant/components/rainmachine/sensor.py b/homeassistant/components/rainmachine/sensor.py index 797420b460f..e2e602b945b 100644 --- a/homeassistant/components/rainmachine/sensor.py +++ b/homeassistant/components/rainmachine/sensor.py @@ -5,8 +5,6 @@ from dataclasses import dataclass from datetime import datetime, timedelta from typing import Any, cast -from regenmaschine.controller import Controller - from homeassistant.components.sensor import ( RestoreSensor, SensorDeviceClass, @@ -17,15 +15,12 @@ from homeassistant.components.sensor import ( from homeassistant.config_entries import ConfigEntry from homeassistant.const import TEMP_CELSIUS, VOLUME_CUBIC_METERS from homeassistant.core import HomeAssistant, callback -from homeassistant.helpers.entity import EntityCategory, EntityDescription +from homeassistant.helpers.entity import EntityCategory from homeassistant.helpers.entity_platform import AddEntitiesCallback -from homeassistant.helpers.update_coordinator import DataUpdateCoordinator from homeassistant.util.dt import utcnow -from . import RainMachineEntity +from . import RainMachineData, RainMachineEntity from .const import ( - DATA_CONTROLLER, - DATA_COORDINATOR, DATA_PROGRAMS, DATA_PROVISION_SETTINGS, DATA_RESTRICTIONS_UNIVERSAL, @@ -33,8 +28,9 @@ from .const import ( DOMAIN, ) from .model import ( - RainMachineDescriptionMixinApiCategory, - RainMachineDescriptionMixinUid, + RainMachineEntityDescription, + RainMachineEntityDescriptionMixinDataKey, + RainMachineEntityDescriptionMixinUid, ) from .util import RUN_STATE_MAP, RunStates, key_exists @@ -50,21 +46,25 @@ TYPE_ZONE_RUN_COMPLETION_TIME = "zone_run_completion_time" @dataclass -class RainMachineSensorDescriptionApiCategory( - SensorEntityDescription, RainMachineDescriptionMixinApiCategory +class RainMachineSensorDataDescription( + SensorEntityDescription, + RainMachineEntityDescription, + RainMachineEntityDescriptionMixinDataKey, ): """Describe a RainMachine sensor.""" @dataclass -class RainMachineSensorDescriptionUid( - SensorEntityDescription, RainMachineDescriptionMixinUid +class RainMachineSensorCompletionTimerDescription( + SensorEntityDescription, + RainMachineEntityDescription, + RainMachineEntityDescriptionMixinUid, ): """Describe a RainMachine sensor.""" SENSOR_DESCRIPTIONS = ( - RainMachineSensorDescriptionApiCategory( + RainMachineSensorDataDescription( key=TYPE_FLOW_SENSOR_CLICK_M3, name="Flow sensor clicks per cubic meter", icon="mdi:water-pump", @@ -75,7 +75,7 @@ SENSOR_DESCRIPTIONS = ( api_category=DATA_PROVISION_SETTINGS, data_key="flowSensorClicksPerCubicMeter", ), - RainMachineSensorDescriptionApiCategory( + RainMachineSensorDataDescription( key=TYPE_FLOW_SENSOR_CONSUMED_LITERS, name="Flow sensor consumed liters", icon="mdi:water-pump", @@ -86,7 +86,7 @@ SENSOR_DESCRIPTIONS = ( api_category=DATA_PROVISION_SETTINGS, data_key="flowSensorWateringClicks", ), - RainMachineSensorDescriptionApiCategory( + RainMachineSensorDataDescription( key=TYPE_FLOW_SENSOR_START_INDEX, name="Flow sensor start index", icon="mdi:water-pump", @@ -96,7 +96,7 @@ SENSOR_DESCRIPTIONS = ( api_category=DATA_PROVISION_SETTINGS, data_key="flowSensorStartIndex", ), - RainMachineSensorDescriptionApiCategory( + RainMachineSensorDataDescription( key=TYPE_FLOW_SENSOR_WATERING_CLICKS, name="Flow sensor clicks", icon="mdi:water-pump", @@ -107,7 +107,7 @@ SENSOR_DESCRIPTIONS = ( api_category=DATA_PROVISION_SETTINGS, data_key="flowSensorWateringClicks", ), - RainMachineSensorDescriptionApiCategory( + RainMachineSensorDataDescription( key=TYPE_FREEZE_TEMP, name="Freeze protect temperature", icon="mdi:thermometer", @@ -125,8 +125,7 @@ async def async_setup_entry( hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback ) -> None: """Set up RainMachine sensors based on a config entry.""" - controller = hass.data[DOMAIN][entry.entry_id][DATA_CONTROLLER] - coordinators = hass.data[DOMAIN][entry.entry_id][DATA_COORDINATOR] + data: RainMachineData = hass.data[DOMAIN][entry.entry_id] api_category_sensor_map = { DATA_PROVISION_SETTINGS: ProvisionSettingsSensor, @@ -134,32 +133,29 @@ async def async_setup_entry( } sensors = [ - api_category_sensor_map[description.api_category]( - entry, coordinator, controller, description - ) + api_category_sensor_map[description.api_category](entry, data, description) for description in SENSOR_DESCRIPTIONS if ( - (coordinator := coordinators[description.api_category]) is not None + (coordinator := data.coordinators[description.api_category]) is not None and coordinator.data and key_exists(coordinator.data, description.data_key) ) ] - program_coordinator = coordinators[DATA_PROGRAMS] - zone_coordinator = coordinators[DATA_ZONES] + program_coordinator = data.coordinators[DATA_PROGRAMS] + zone_coordinator = data.coordinators[DATA_ZONES] for uid, program in program_coordinator.data.items(): sensors.append( ProgramTimeRemainingSensor( entry, - program_coordinator, - zone_coordinator, - controller, - RainMachineSensorDescriptionUid( + data, + RainMachineSensorCompletionTimerDescription( key=f"{TYPE_PROGRAM_RUN_COMPLETION_TIME}_{uid}", name=f"{program['name']} Run Completion Time", device_class=SensorDeviceClass.TIMESTAMP, entity_category=EntityCategory.DIAGNOSTIC, + api_category=DATA_PROGRAMS, uid=uid, ), ) @@ -169,13 +165,13 @@ async def async_setup_entry( sensors.append( ZoneTimeRemainingSensor( entry, - zone_coordinator, - controller, - RainMachineSensorDescriptionUid( + data, + RainMachineSensorCompletionTimerDescription( key=f"{TYPE_ZONE_RUN_COMPLETION_TIME}_{uid}", name=f"{zone['name']} Run Completion Time", device_class=SensorDeviceClass.TIMESTAMP, entity_category=EntityCategory.DIAGNOSTIC, + api_category=DATA_ZONES, uid=uid, ), ) @@ -187,17 +183,16 @@ async def async_setup_entry( class TimeRemainingSensor(RainMachineEntity, RestoreSensor): """Define a sensor that shows the amount of time remaining for an activity.""" - entity_description: RainMachineSensorDescriptionUid + entity_description: RainMachineSensorCompletionTimerDescription def __init__( self, entry: ConfigEntry, - coordinator: DataUpdateCoordinator, - controller: Controller, - description: EntityDescription, + data: RainMachineData, + description: RainMachineSensorCompletionTimerDescription, ) -> None: """Initialize.""" - super().__init__(entry, coordinator, controller, description) + super().__init__(entry, data, description) self._current_run_state: RunStates | None = None self._previous_run_state: RunStates | None = None @@ -256,19 +251,6 @@ class TimeRemainingSensor(RainMachineEntity, RestoreSensor): class ProgramTimeRemainingSensor(TimeRemainingSensor): """Define a sensor that shows the amount of time remaining for a program.""" - def __init__( - self, - entry: ConfigEntry, - program_coordinator: DataUpdateCoordinator, - zone_coordinator: DataUpdateCoordinator, - controller: Controller, - description: EntityDescription, - ) -> None: - """Initialize.""" - super().__init__(entry, program_coordinator, controller, description) - - self._zone_coordinator = zone_coordinator - @property def status_key(self) -> str: """Return the data key that contains the activity status.""" @@ -277,7 +259,7 @@ class ProgramTimeRemainingSensor(TimeRemainingSensor): def calculate_seconds_remaining(self) -> int: """Calculate the number of seconds remaining.""" return sum( - self._zone_coordinator.data[zone["id"]]["remaining"] + self._data.coordinators[DATA_ZONES].data[zone["id"]]["remaining"] for zone in [z for z in self.activity_data["wateringTimes"] if z["active"]] ) @@ -285,6 +267,8 @@ class ProgramTimeRemainingSensor(TimeRemainingSensor): class ProvisionSettingsSensor(RainMachineEntity, SensorEntity): """Define a sensor that handles provisioning data.""" + entity_description: RainMachineSensorDataDescription + @callback def update_from_latest_data(self) -> None: """Update the state.""" @@ -315,6 +299,8 @@ class ProvisionSettingsSensor(RainMachineEntity, SensorEntity): class UniversalRestrictionsSensor(RainMachineEntity, SensorEntity): """Define a sensor that handles universal restrictions data.""" + entity_description: RainMachineSensorDataDescription + @callback def update_from_latest_data(self) -> None: """Update the state.""" diff --git a/homeassistant/components/rainmachine/switch.py b/homeassistant/components/rainmachine/switch.py index 1fcdab49836..ee6ac670840 100644 --- a/homeassistant/components/rainmachine/switch.py +++ b/homeassistant/components/rainmachine/switch.py @@ -7,7 +7,6 @@ from dataclasses import dataclass from datetime import datetime from typing import Any -from regenmaschine.controller import Controller from regenmaschine.errors import RequestError import voluptuous as vol @@ -19,19 +18,16 @@ from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers import config_validation as cv, entity_platform from homeassistant.helpers.entity import EntityCategory from homeassistant.helpers.entity_platform import AddEntitiesCallback -from homeassistant.helpers.update_coordinator import DataUpdateCoordinator -from . import RainMachineEntity, async_update_programs_and_zones +from . import RainMachineData, RainMachineEntity, async_update_programs_and_zones from .const import ( CONF_ZONE_RUN_TIME, - DATA_CONTROLLER, - DATA_COORDINATOR, DATA_PROGRAMS, DATA_ZONES, DEFAULT_ZONE_RUN, DOMAIN, ) -from .model import RainMachineDescriptionMixinUid +from .model import RainMachineEntityDescription, RainMachineEntityDescriptionMixinUid from .util import RUN_STATE_MAP ATTR_AREA = "area" @@ -110,7 +106,9 @@ VEGETATION_MAP = { @dataclass class RainMachineSwitchDescription( - SwitchEntityDescription, RainMachineDescriptionMixinUid + SwitchEntityDescription, + RainMachineEntityDescription, + RainMachineEntityDescriptionMixinUid, ): """Describe a RainMachine switch.""" @@ -137,30 +135,27 @@ async def async_setup_entry( ): platform.async_register_entity_service(service_name, schema, method) - data = hass.data[DOMAIN][entry.entry_id] - controller = data[DATA_CONTROLLER] - program_coordinator = data[DATA_COORDINATOR][DATA_PROGRAMS] - zone_coordinator = data[DATA_COORDINATOR][DATA_ZONES] + data: RainMachineData = hass.data[DOMAIN][entry.entry_id] entities: list[RainMachineActivitySwitch | RainMachineEnabledSwitch] = [] - - for kind, coordinator, switch_class, switch_enabled_class in ( - ("program", program_coordinator, RainMachineProgram, RainMachineProgramEnabled), - ("zone", zone_coordinator, RainMachineZone, RainMachineZoneEnabled), + for kind, api_category, switch_class, switch_enabled_class in ( + ("program", DATA_PROGRAMS, RainMachineProgram, RainMachineProgramEnabled), + ("zone", DATA_ZONES, RainMachineZone, RainMachineZoneEnabled), ): - for uid, data in coordinator.data.items(): - name = data["name"].capitalize() + coordinator = data.coordinators[api_category] + for uid, activity in coordinator.data.items(): + name = activity["name"].capitalize() # Add a switch to start/stop the program or zone: entities.append( switch_class( entry, - coordinator, - controller, + data, RainMachineSwitchDescription( key=f"{kind}_{uid}", name=name, icon="mdi:water", + api_category=api_category, uid=uid, ), ) @@ -170,13 +165,13 @@ async def async_setup_entry( entities.append( switch_enabled_class( entry, - coordinator, - controller, + data, RainMachineSwitchDescription( key=f"{kind}_{uid}_enabled", name=f"{name} enabled", entity_category=EntityCategory.CONFIG, icon="mdi:cog", + api_category=api_category, uid=uid, ), ) @@ -193,12 +188,11 @@ class RainMachineBaseSwitch(RainMachineEntity, SwitchEntity): def __init__( self, entry: ConfigEntry, - coordinator: DataUpdateCoordinator, - controller: Controller, + data: RainMachineData, description: RainMachineSwitchDescription, ) -> None: """Initialize.""" - super().__init__(entry, coordinator, controller, description) + super().__init__(entry, data, description) self._attr_is_on = False self._entry = entry @@ -299,13 +293,13 @@ class RainMachineProgram(RainMachineActivitySwitch): async def async_turn_off_when_active(self, **kwargs: Any) -> None: """Turn the switch off when its associated activity is active.""" await self._async_run_api_coroutine( - self._controller.programs.stop(self.entity_description.uid) + self._data.controller.programs.stop(self.entity_description.uid) ) async def async_turn_on_when_active(self, **kwargs: Any) -> None: """Turn the switch on when its associated activity is active.""" await self._async_run_api_coroutine( - self._controller.programs.start(self.entity_description.uid) + self._data.controller.programs.start(self.entity_description.uid) ) @callback @@ -342,10 +336,10 @@ class RainMachineProgramEnabled(RainMachineEnabledSwitch): """Disable the program.""" tasks = [ self._async_run_api_coroutine( - self._controller.programs.stop(self.entity_description.uid) + self._data.controller.programs.stop(self.entity_description.uid) ), self._async_run_api_coroutine( - self._controller.programs.disable(self.entity_description.uid) + self._data.controller.programs.disable(self.entity_description.uid) ), ] @@ -354,7 +348,7 @@ class RainMachineProgramEnabled(RainMachineEnabledSwitch): async def async_turn_on(self, **kwargs: Any) -> None: """Enable the program.""" await self._async_run_api_coroutine( - self._controller.programs.enable(self.entity_description.uid) + self._data.controller.programs.enable(self.entity_description.uid) ) @@ -372,13 +366,13 @@ class RainMachineZone(RainMachineActivitySwitch): async def async_turn_off_when_active(self, **kwargs: Any) -> None: """Turn the switch off when its associated activity is active.""" await self._async_run_api_coroutine( - self._controller.zones.stop(self.entity_description.uid) + self._data.controller.zones.stop(self.entity_description.uid) ) async def async_turn_on_when_active(self, **kwargs: Any) -> None: """Turn the switch on when its associated activity is active.""" await self._async_run_api_coroutine( - self._controller.zones.start( + self._data.controller.zones.start( self.entity_description.uid, kwargs.get("duration", self._entry.options[CONF_ZONE_RUN_TIME]), ) @@ -426,10 +420,10 @@ class RainMachineZoneEnabled(RainMachineEnabledSwitch): """Disable the zone.""" tasks = [ self._async_run_api_coroutine( - self._controller.zones.stop(self.entity_description.uid) + self._data.controller.zones.stop(self.entity_description.uid) ), self._async_run_api_coroutine( - self._controller.zones.disable(self.entity_description.uid) + self._data.controller.zones.disable(self.entity_description.uid) ), ] @@ -438,5 +432,5 @@ class RainMachineZoneEnabled(RainMachineEnabledSwitch): async def async_turn_on(self, **kwargs: Any) -> None: """Enable the zone.""" await self._async_run_api_coroutine( - self._controller.zones.enable(self.entity_description.uid) + self._data.controller.zones.enable(self.entity_description.uid) )