Add a calendar entity to ReCollect Waste (home-assistant#85347)

* Add a calendar entity to ReCollect Waste

* Simplify and ensure return None

* Ensure end date is after start date

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}"

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 + datetime.timedelta(days=1),
+    )
+
+
+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:
+            self._event = None
+        else:
+            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
+        ]

homeassistant/components/recollect_waste/sensor.py

@@ -1,21 +1,23 @@
 """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,
     SensorEntity,
     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:
@@ -67,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]],
@@ -82,25 +76,19 @@ def __init__(
     @callback
     def _handle_coordinator_update(self) -> None:
         """Handle updated data from the coordinator."""
-        if self.entity_description.key == SENSOR_TYPE_CURRENT_PICKUP:
-            try:
-                event = self.coordinator.data[0]
-            except IndexError:
-                LOGGER.error("No current pickup found")
-                return
+        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:
+            for _ in range(pickup_index):
+                event = next(relevant_events)
+        except StopIteration:
+            LOGGER.info("No pickup event found for %s", self.entity_description.key)
+            self._attr_extra_state_attributes = {}
+            self._attr_native_value = None
         else:
-            try:
-                event = self.coordinator.data[1]
-            except IndexError:
-                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._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

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
+    ]