HomeAssistant_Config/custom_components/waste_collection_schedule/calendar.py

"""Calendar platform support for Waste Collection Schedule."""

import logging
from datetime import datetime, timedelta

from homeassistant.components.calendar import CalendarEntity, CalendarEvent
from homeassistant.core import HomeAssistant

from custom_components.waste_collection_schedule.waste_collection_schedule.scraper import Scraper

_LOGGER = logging.getLogger(__name__)


async def async_setup_platform(hass, config, async_add_entities, discovery_info=None):
    """Set up calendar platform."""
    # We only want this platform to be set up via discovery.
    if discovery_info is None:
        return

    entities = []

    api = discovery_info["api"]

    for scraper in api.scrapers:
        dedicated_calendar_types = scraper.get_dedicated_calendar_types()
        global_calendar_types = scraper.get_global_calendar_types()

        if dedicated_calendar_types is not None:
            for type in dedicated_calendar_types:
                unique_id = calc_unique_calendar_id(scraper, type)

                entities.append(
                    WasteCollectionCalendar(
                        api,
                        scraper,
                        scraper.get_calendar_title_for_type(type),
                        [scraper.get_collection_type(type)],
                        unique_id,
                    )
                )

        if global_calendar_types is not None or dedicated_calendar_types is None:
            unique_id = calc_unique_calendar_id(scraper)
            entities.append(
                WasteCollectionCalendar(
                    api,
                    scraper,
                    scraper.calendar_title,
                    [
                        scraper.get_collection_type(type)
                        for type in global_calendar_types
                    ]
                    if global_calendar_types is not None
                    else None,
                    unique_id,
                )
            )

    async_add_entities(entities)


class WasteCollectionCalendar(CalendarEntity):
    """Calendar entity class."""

    def __init__(self, api, scraper, name, types, unique_id: str):
        self._api = api
        self._scraper = scraper
        self._name = name
        self._types = types
        self._unique_id = unique_id
        self._attr_unique_id = unique_id

    @property
    def name(self):
        """Return entity name."""
        return self._name

    @property
    def event(self):
        """Return next collection event."""
        collections = self._scraper.get_upcoming(
            count=1, include_today=True, types=self._types
        )

        if len(collections) == 0:
            return None
        else:
            return self._convert(collections[0])

    async def async_get_events(
        self, hass: HomeAssistant, start_date: datetime, end_date: datetime
    ):
        """Return all events within specified time span."""
        events = []

        for collection in self._scraper.get_upcoming(
            include_today=True, types=self._types
        ):
            event = self._convert(collection)

            if start_date <= event.start_datetime_local <= end_date:
                events.append(event)

        return events

    def _convert(self, collection) -> CalendarEvent:
        """Convert an collection into a Home Assistant calendar event."""
        return CalendarEvent(
            summary=collection.type,
            start=collection.date,
            end=collection.date + timedelta(days=1),
        )


def calc_unique_calendar_id(scraper: Scraper, type: str = None):
    return scraper.unique_id + ("_" + type if type is not None else "") + "_calendar"