From 9152e34cf367998209cfda90608cafb6042f4800 Mon Sep 17 00:00:00 2001 From: Aaron Bach Date: Fri, 6 Jan 2023 15:55:41 -0700 Subject: [PATCH 1/3] Add a calendar entity to ReCollect Waste --- .../components/recollect_waste/__init__.py | 8 +- .../components/recollect_waste/calendar.py | 96 +++++++++++++++++++ .../components/recollect_waste/sensor.py | 63 ++++++------ .../components/recollect_waste/util.py | 19 ++++ 4 files changed, 150 insertions(+), 36 deletions(-) create mode 100644 homeassistant/components/recollect_waste/calendar.py create mode 100644 homeassistant/components/recollect_waste/util.py diff --git a/homeassistant/components/recollect_waste/__init__.py b/homeassistant/components/recollect_waste/__init__.py index 0d7f35b6e62cee..21cf574d548b69 100644 --- a/homeassistant/components/recollect_waste/__init__.py +++ b/homeassistant/components/recollect_waste/__init__.py @@ -1,7 +1,7 @@ """The ReCollect Waste integration.""" from __future__ import annotations -from datetime import date, timedelta +from datetime import timedelta from typing import Any from aiorecollect.client import Client, PickupEvent @@ -18,7 +18,7 @@ DEFAULT_NAME = "recollect_waste" DEFAULT_UPDATE_INTERVAL = timedelta(days=1) -PLATFORMS = [Platform.SENSOR] +PLATFORMS = [Platform.CALENDAR, Platform.SENSOR] async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: @@ -31,9 +31,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: async def async_get_pickup_events() -> list[PickupEvent]: """Get the next pickup.""" try: - return await client.async_get_pickup_events( - start_date=date.today(), end_date=date.today() + timedelta(weeks=4) - ) + return await client.async_get_pickup_events() except RecollectError as err: raise UpdateFailed( f"Error while requesting data from ReCollect: {err}" diff --git a/homeassistant/components/recollect_waste/calendar.py b/homeassistant/components/recollect_waste/calendar.py new file mode 100644 index 00000000000000..bb6ef4089e1bf0 --- /dev/null +++ b/homeassistant/components/recollect_waste/calendar.py @@ -0,0 +1,96 @@ +"""Support for ReCollect Waste calendars.""" +from __future__ import annotations + +import datetime + +from aiorecollect.client import PickupEvent + +from homeassistant.components.calendar import CalendarEntity, CalendarEvent +from homeassistant.config_entries import ConfigEntry +from homeassistant.core import HomeAssistant, callback +from homeassistant.helpers.entity_platform import AddEntitiesCallback +from homeassistant.helpers.update_coordinator import DataUpdateCoordinator + +from .const import DOMAIN +from .entity import ReCollectWasteEntity +from .util import async_get_pickup_type_names + + +@callback +def async_get_calendar_event_from_pickup_event( + entry: ConfigEntry, pickup_event: PickupEvent +) -> CalendarEvent: + """Get a HASS CalendarEvent from an aiorecollect PickupEvent.""" + pickup_type_string = ", ".join( + async_get_pickup_type_names(entry, pickup_event.pickup_types) + ) + return CalendarEvent( + summary="ReCollect Waste Pickup", + description=f"Pickup types: {pickup_type_string}", + location=pickup_event.area_name, + start=pickup_event.date, + end=pickup_event.date, + ) + + +async def async_setup_entry( + hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback +) -> None: + """Set up ReCollect Waste sensors based on a config entry.""" + coordinator: DataUpdateCoordinator[list[PickupEvent]] = hass.data[DOMAIN][ + entry.entry_id + ] + + async_add_entities([ReCollectWasteCalendar(coordinator, entry)]) + + +class ReCollectWasteCalendar(ReCollectWasteEntity, CalendarEntity): + """Define a ReCollect Waste calendar.""" + + _attr_icon = "mdi:delete-empty" + + def __init__( + self, + coordinator: DataUpdateCoordinator[list[PickupEvent]], + entry: ConfigEntry, + ) -> None: + """Initialize the ReCollect Waste entity.""" + super().__init__(coordinator, entry) + + self._attr_unique_id = f"{self._identifier}_calendar" + self._event: CalendarEvent | None = None + + @property + def event(self) -> CalendarEvent | None: + """Return the next upcoming event.""" + return self._event + + @callback + def _handle_coordinator_update(self) -> None: + """Handle updated data from the coordinator.""" + try: + current_event = next( + event + for event in self.coordinator.data + if event.date >= datetime.date.today() + ) + except StopIteration: + return + + self._event = async_get_calendar_event_from_pickup_event( + self._entry, current_event + ) + + super()._handle_coordinator_update() + + async def async_get_events( + self, + hass: HomeAssistant, + start_date: datetime.datetime, + end_date: datetime.datetime, + ) -> list[CalendarEvent]: + """Return calendar events within a datetime range.""" + return [ + async_get_calendar_event_from_pickup_event(self._entry, event) + for event in self.coordinator.data + ] diff --git a/homeassistant/components/recollect_waste/sensor.py b/homeassistant/components/recollect_waste/sensor.py index eff3ce0b9a3521..bad50f6bab8ebe 100644 --- a/homeassistant/components/recollect_waste/sensor.py +++ b/homeassistant/components/recollect_waste/sensor.py @@ -1,7 +1,9 @@ """Support for ReCollect Waste sensors.""" from __future__ import annotations -from aiorecollect.client import PickupEvent, PickupType +from datetime import date + +from aiorecollect.client import PickupEvent from homeassistant.components.sensor import ( SensorDeviceClass, @@ -9,13 +11,13 @@ SensorEntityDescription, ) from homeassistant.config_entries import ConfigEntry -from homeassistant.const import CONF_FRIENDLY_NAME from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.update_coordinator import DataUpdateCoordinator from .const import DOMAIN, LOGGER from .entity import ReCollectWasteEntity +from .util import async_get_pickup_type_names ATTR_PICKUP_TYPES = "pickup_types" ATTR_AREA_NAME = "area_name" @@ -35,19 +37,6 @@ ) -@callback -def async_get_pickup_type_names( - entry: ConfigEntry, pickup_types: list[PickupType] -) -> list[str]: - """Return proper pickup type names from their associated objects.""" - return [ - t.friendly_name - if entry.options.get(CONF_FRIENDLY_NAME) and t.friendly_name - else t.name - for t in pickup_types - ] - - async def async_setup_entry( hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback ) -> None: @@ -79,28 +68,40 @@ def __init__( self._attr_unique_id = f"{self._identifier}_{description.key}" self.entity_description = description + @callback + def _async_write_state_from_event(self, event: PickupEvent) -> None: + """Write the entity state from a pickup event.""" + pickup_types = async_get_pickup_type_names(self._entry, event.pickup_types) + self._attr_extra_state_attributes.update( + { + ATTR_AREA_NAME: event.area_name, + ATTR_PICKUP_TYPES: pickup_types, + } + ) + self._attr_native_value = event.date + @callback def _handle_coordinator_update(self) -> None: """Handle updated data from the coordinator.""" + try: + current_event = next( + event for event in self.coordinator.data if event.date >= date.today() + ) + except StopIteration: + LOGGER.error("No current pickup found") + return + if self.entity_description.key == SENSOR_TYPE_CURRENT_PICKUP: - try: - event = self.coordinator.data[0] - except IndexError: - LOGGER.error("No current pickup found") - return + self._async_write_state_from_event(current_event) else: try: - event = self.coordinator.data[1] - except IndexError: + next_event = next( + event + for event in self.coordinator.data + if event.date > current_event.date + ) + except StopIteration: LOGGER.info("No next pickup found") return - self._attr_extra_state_attributes.update( - { - ATTR_PICKUP_TYPES: async_get_pickup_type_names( - self._entry, event.pickup_types - ), - ATTR_AREA_NAME: event.area_name, - } - ) - self._attr_native_value = event.date + self._async_write_state_from_event(next_event) diff --git a/homeassistant/components/recollect_waste/util.py b/homeassistant/components/recollect_waste/util.py new file mode 100644 index 00000000000000..185078f297c139 --- /dev/null +++ b/homeassistant/components/recollect_waste/util.py @@ -0,0 +1,19 @@ +"""Define ReCollect Waste utilities.""" +from aiorecollect.client import PickupType + +from homeassistant.config_entries import ConfigEntry +from homeassistant.const import CONF_FRIENDLY_NAME +from homeassistant.core import callback + + +@callback +def async_get_pickup_type_names( + entry: ConfigEntry, pickup_types: list[PickupType] +) -> list[str]: + """Return proper pickup type names from their associated objects.""" + return [ + t.friendly_name + if entry.options.get(CONF_FRIENDLY_NAME) and t.friendly_name + else t.name + for t in pickup_types + ] From d91510b2450b6fc3f4591641bb47e6a111958195 Mon Sep 17 00:00:00 2001 From: Aaron Bach Date: Sat, 7 Jan 2023 13:00:55 -0700 Subject: [PATCH 2/3] Simplify and ensure return None --- .../components/recollect_waste/calendar.py | 10 ++-- .../components/recollect_waste/sensor.py | 49 +++++++------------ 2 files changed, 23 insertions(+), 36 deletions(-) diff --git a/homeassistant/components/recollect_waste/calendar.py b/homeassistant/components/recollect_waste/calendar.py index bb6ef4089e1bf0..54c4c9fcdbf005 100644 --- a/homeassistant/components/recollect_waste/calendar.py +++ b/homeassistant/components/recollect_waste/calendar.py @@ -75,11 +75,11 @@ def _handle_coordinator_update(self) -> None: if event.date >= datetime.date.today() ) except StopIteration: - return - - self._event = async_get_calendar_event_from_pickup_event( - self._entry, current_event - ) + self._event = None + else: + self._event = async_get_calendar_event_from_pickup_event( + self._entry, current_event + ) super()._handle_coordinator_update() diff --git a/homeassistant/components/recollect_waste/sensor.py b/homeassistant/components/recollect_waste/sensor.py index bad50f6bab8ebe..1261a6a9b5c18c 100644 --- a/homeassistant/components/recollect_waste/sensor.py +++ b/homeassistant/components/recollect_waste/sensor.py @@ -56,6 +56,11 @@ class ReCollectWasteSensor(ReCollectWasteEntity, SensorEntity): _attr_device_class = SensorDeviceClass.DATE + PICKUP_INDEX_MAP = { + SENSOR_TYPE_CURRENT_PICKUP: 1, + SENSOR_TYPE_NEXT_PICKUP: 2, + } + def __init__( self, coordinator: DataUpdateCoordinator[list[PickupEvent]], @@ -68,40 +73,22 @@ def __init__( self._attr_unique_id = f"{self._identifier}_{description.key}" self.entity_description = description - @callback - def _async_write_state_from_event(self, event: PickupEvent) -> None: - """Write the entity state from a pickup event.""" - pickup_types = async_get_pickup_type_names(self._entry, event.pickup_types) - self._attr_extra_state_attributes.update( - { - ATTR_AREA_NAME: event.area_name, - ATTR_PICKUP_TYPES: pickup_types, - } - ) - self._attr_native_value = event.date - @callback def _handle_coordinator_update(self) -> None: """Handle updated data from the coordinator.""" + relevant_events = (e for e in self.coordinator.data if e.date >= date.today()) + pickup_index = self.PICKUP_INDEX_MAP[self.entity_description.key] + try: - current_event = next( - event for event in self.coordinator.data if event.date >= date.today() - ) + for _ in range(pickup_index): + event = next(relevant_events) except StopIteration: - LOGGER.error("No current pickup found") - return - - if self.entity_description.key == SENSOR_TYPE_CURRENT_PICKUP: - self._async_write_state_from_event(current_event) + LOGGER.info("No pickup event found for %s", self.entity_description.key) + self._attr_extra_state_attributes = {} + self._attr_native_value = None else: - try: - next_event = next( - event - for event in self.coordinator.data - if event.date > current_event.date - ) - except StopIteration: - LOGGER.info("No next pickup found") - return - - self._async_write_state_from_event(next_event) + self._attr_extra_state_attributes[ATTR_AREA_NAME] = event.area_name + self._attr_extra_state_attributes[ + ATTR_PICKUP_TYPES + ] = async_get_pickup_type_names(self._entry, event.pickup_types) + self._attr_native_value = event.date From b2ef0e65a939f6a7ab1da295ff8c717aa91f8413 Mon Sep 17 00:00:00 2001 From: Aaron Bach Date: Mon, 9 Jan 2023 09:44:51 -0700 Subject: [PATCH 3/3] Ensure end date is after start date --- homeassistant/components/recollect_waste/calendar.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/homeassistant/components/recollect_waste/calendar.py b/homeassistant/components/recollect_waste/calendar.py index 54c4c9fcdbf005..b1af5885dd3498 100644 --- a/homeassistant/components/recollect_waste/calendar.py +++ b/homeassistant/components/recollect_waste/calendar.py @@ -29,7 +29,7 @@ def async_get_calendar_event_from_pickup_event( description=f"Pickup types: {pickup_type_string}", location=pickup_event.area_name, start=pickup_event.date, - end=pickup_event.date, + end=pickup_event.date + datetime.timedelta(days=1), )