Update caldav to use an DataUpdateCoordinator for fetching data (#102089)

homeassistant/components/caldav/calendar.py

@@ -24,12 +24,16 @@ from homeassistant.const import (
     CONF_USERNAME,
     CONF_VERIFY_SSL,
 )
-from homeassistant.core import HomeAssistant
+from homeassistant.core import HomeAssistant, callback
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.entity import generate_entity_id
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType
-from homeassistant.util import Throttle, dt as dt_util
+from homeassistant.helpers.update_coordinator import (
+    CoordinatorEntity,
+    DataUpdateCoordinator,
+)
+from homeassistant.util import dt as dt_util
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -103,16 +107,14 @@ def setup_platform(
             name = cust_calendar[CONF_NAME]
             device_id = f"{cust_calendar[CONF_CALENDAR]} {cust_calendar[CONF_NAME]}"
             entity_id = generate_entity_id(ENTITY_ID_FORMAT, device_id, hass=hass)
-            calendar_devices.append(
+            coordinator = CalDavUpdateCoordinator(
-                WebDavCalendarEntity(
+                hass,
-                    name=name,
                 calendar=calendar,
-                    entity_id=entity_id,
                 days=days,
-                    all_day=True,
+                include_all_day=True,
                 search=cust_calendar[CONF_SEARCH],
             )
-            )
+            calendar_devices.append(WebDavCalendarEntity(name, entity_id, coordinator))
 
         # Create a default calendar if there was no custom one for all calendars
         # that support events.
@@ -130,24 +132,26 @@ def setup_platform(
             name = calendar.name
             device_id = calendar.name
             entity_id = generate_entity_id(ENTITY_ID_FORMAT, device_id, hass=hass)
-            calendar_devices.append(
+            coordinator = CalDavUpdateCoordinator(
-                WebDavCalendarEntity(name, calendar, entity_id, days)
+                hass,
+                calendar=calendar,
+                days=days,
+                include_all_day=False,
+                search=None,
             )
+            calendar_devices.append(WebDavCalendarEntity(name, entity_id, coordinator))
 
     add_entities(calendar_devices, True)
 
 
-class WebDavCalendarEntity(CalendarEntity):
+class WebDavCalendarEntity(
+    CoordinatorEntity["CalDavUpdateCoordinator"], CalendarEntity
+):
     """A device for getting the next Task from a WebDav Calendar."""
 
-    def __init__(self, name, calendar, entity_id, days, all_day=False, search=None):
+    def __init__(self, name, entity_id, coordinator):
         """Create the WebDav Calendar Event Device."""
-        self.data = WebDavCalendarData(
+        super().__init__(coordinator)
-            calendar=calendar,
-            days=days,
-            include_all_day=all_day,
-            search=search,
-        )
         self.entity_id = entity_id
         self._event: CalendarEvent | None = None
         self._attr_name = name
@@ -161,31 +165,42 @@ class WebDavCalendarEntity(CalendarEntity):
         self, hass: HomeAssistant, start_date: datetime, end_date: datetime
     ) -> list[CalendarEvent]:
         """Get all events in a specific time frame."""
-        return await self.data.async_get_events(hass, start_date, end_date)
+        return await self.coordinator.async_get_events(hass, start_date, end_date)
 
-    def update(self) -> None:
+    @callback
+    def _handle_coordinator_update(self) -> None:
         """Update event data."""
-        self.data.update()
+        self._event = self.coordinator.data
-        self._event = self.data.event
         self._attr_extra_state_attributes = {
             "offset_reached": is_offset_reached(
-                self._event.start_datetime_local, self.data.offset
+                self._event.start_datetime_local, self.coordinator.offset
             )
             if self._event
             else False
         }
+        super()._handle_coordinator_update()
+
+    async def async_added_to_hass(self) -> None:
+        """When entity is added to hass update state from existing coordinator data."""
+        await super().async_added_to_hass()
+        self._handle_coordinator_update()
 
 
-class WebDavCalendarData:
+class CalDavUpdateCoordinator(DataUpdateCoordinator[CalendarEvent | None]):
     """Class to utilize the calendar dav client object to get next event."""
 
-    def __init__(self, calendar, days, include_all_day, search):
+    def __init__(self, hass, calendar, days, include_all_day, search):
         """Set up how we are going to search the WebDav calendar."""
+        super().__init__(
+            hass,
+            _LOGGER,
+            name=f"CalDAV {calendar.name}",
+            update_interval=MIN_TIME_BETWEEN_UPDATES,
+        )
         self.calendar = calendar
         self.days = days
         self.include_all_day = include_all_day
         self.search = search
-        self.event = None
         self.offset = None
 
     async def async_get_events(
@@ -222,19 +237,21 @@ class WebDavCalendarData:
 
         return event_list
 
-    @Throttle(MIN_TIME_BETWEEN_UPDATES)
+    async def _async_update_data(self) -> CalendarEvent | None:
-    def update(self):
         """Get the latest data."""
         start_of_today = dt_util.start_of_local_day()
         start_of_tomorrow = dt_util.start_of_local_day() + timedelta(days=self.days)
 
         # We have to retrieve the results for the whole day as the server
         # won't return events that have already started
-        results = self.calendar.search(
+        results = await self.hass.async_add_executor_job(
+            partial(
+                self.calendar.search,
                 start=start_of_today,
                 end=start_of_tomorrow,
                 event=True,
                 expand=True,
+            ),
         )
 
         # Create new events for each recurrence of an event that happens today.
@@ -247,12 +264,15 @@ class WebDavCalendarData:
                 continue
             vevent = event.instance.vevent
             for start_dt in vevent.getrruleset() or []:
-                _start_of_today = start_of_today
+                _start_of_today: date | datetime
-                _start_of_tomorrow = start_of_tomorrow
+                _start_of_tomorrow: datetime | date
                 if self.is_all_day(vevent):
                     start_dt = start_dt.date()
-                    _start_of_today = _start_of_today.date()
+                    _start_of_today = start_of_today.date()
-                    _start_of_tomorrow = _start_of_tomorrow.date()
+                    _start_of_tomorrow = start_of_tomorrow.date()
+                else:
+                    _start_of_today = start_of_today
+                    _start_of_tomorrow = start_of_tomorrow
                 if _start_of_today <= start_dt < _start_of_tomorrow:
                     new_event = event.copy()
                     new_vevent = new_event.instance.vevent
@@ -293,21 +313,21 @@ class WebDavCalendarData:
                 len(vevents),
                 self.calendar.name,
             )
-            self.event = None
+            self.offset = None
-            return
+            return None
 
         # Populate the entity attributes with the event values
         (summary, offset) = extract_offset(
             self.get_attr_value(vevent, "summary") or "", OFFSET
         )
-        self.event = CalendarEvent(
+        self.offset = offset
+        return CalendarEvent(
             summary=summary,
             start=self.to_local(vevent.dtstart.value),
             end=self.to_local(self.get_end_date(vevent)),
             location=self.get_attr_value(vevent, "location"),
             description=self.get_attr_value(vevent, "description"),
         )
-        self.offset = offset
 
     @staticmethod
     def is_matching(vevent, search):
@@ -333,15 +353,15 @@ class WebDavCalendarData:
     @staticmethod
     def is_over(vevent):
         """Return if the event is over."""
-        return dt_util.now() >= WebDavCalendarData.to_datetime(
+        return dt_util.now() >= CalDavUpdateCoordinator.to_datetime(
-            WebDavCalendarData.get_end_date(vevent)
+            CalDavUpdateCoordinator.get_end_date(vevent)
         )
 
     @staticmethod
     def to_datetime(obj):
         """Return a datetime."""
         if isinstance(obj, datetime):
-            return WebDavCalendarData.to_local(obj)
+            return CalDavUpdateCoordinator.to_local(obj)
         return datetime.combine(obj, time.min).replace(tzinfo=dt_util.DEFAULT_TIME_ZONE)
 
     @staticmethod