__init__.py

Go to the documentation of this file.

"""Support for schedules in Home Assistant."""
 
from __future__ import annotations
 
from collections.abc import Callable
from datetime import datetime, time, timedelta
import itertools
from typing import Any, Literal
 
import voluptuous as vol
 
from homeassistant.const import (
    ATTR_EDITABLE,
    CONF_ICON,
    CONF_ID,
    CONF_NAME,
    SERVICE_RELOAD,
    STATE_OFF,
    STATE_ON,
)
from homeassistant.core import HomeAssistant, ServiceCall, callback
from homeassistant.helpers.collection import (
    CollectionEntity,
    DictStorageCollection,
    DictStorageCollectionWebsocket,
    IDManager,
    SerializedStorageCollection,
    YamlCollection,
    sync_entity_lifecycle,
)
import homeassistant.helpers.config_validation as cv
from homeassistant.helpers.entity_component import EntityComponent
from homeassistant.helpers.event import async_track_point_in_utc_time
from homeassistant.helpers.service import async_register_admin_service
from homeassistant.helpers.storage import Store
from homeassistant.helpers.typing import ConfigType, VolDictType
from homeassistant.util import dt as dt_util
 
from .const import (
    ATTR_NEXT_EVENT,
    CONF_ALL_DAYS,
    CONF_DATA,
    CONF_FROM,
    CONF_TO,
    DOMAIN,
    LOGGER,
    WEEKDAY_TO_CONF,
)
 
STORAGE_VERSION = 1
STORAGE_VERSION_MINOR = 1
 
 
def valid_schedule(schedule: list[dict[str, str]]) -> list[dict[str, str]]:
    """Validate the schedule of time ranges.
 
    Ensure they have no overlap and the end time is greater than the start time.
    """
    # Empty schedule is valid
    if not schedule:
        return schedule
 
    # Sort the schedule by start times
    schedule = sorted(schedule, key=lambda time_range: time_range[CONF_FROM])
 
    # Check if the start time of the next event is before the end time of the previous event
    previous_to = None
    for time_range in schedule:
        if time_range[CONF_FROM] >= time_range[CONF_TO]:
            raise vol.Invalid(
                f"Invalid time range, from {time_range[CONF_FROM]} is after"
                f" {time_range[CONF_TO]}"
            )
 
        # Check if the from time of the event is after the to time of the previous event
        if previous_to is not None and previous_to > time_range[CONF_FROM]:  # type: ignore[unreachable]
            raise vol.Invalid("Overlapping times found in schedule")
 
        previous_to = time_range[CONF_TO]
 
    return schedule
 
 
def deserialize_to_time(value: Any) -> Any:
    """Convert 24:00 and 24:00:00 to time.max."""
    if not isinstance(value, str):
        return cv.time(value)
 
    parts = value.split(":")
    if len(parts) < 2:
        return cv.time(value)
    hour = int(parts[0])
    minute = int(parts[1])
 
    if hour == 24 and minute == 0:
        return time.max
 
    return cv.time(value)
 
 
def serialize_to_time(value: Any) -> Any:
    """Convert time.max to 24:00:00."""
    if value == time.max:
        return "24:00:00"
    return vol.Coerce(str)(value)
 
 
BASE_SCHEMA: VolDictType = {
    vol.Required(CONF_NAME): vol.All(str, vol.Length(min=1)),
    vol.Optional(CONF_ICON): cv.icon,
}
 
# Extra data that the user can set on each time range
CUSTOM_DATA_SCHEMA = vol.Schema({str: vol.Any(bool, str, int, float)})
 
TIME_RANGE_SCHEMA: VolDictType = {
    vol.Required(CONF_FROM): cv.time,
    vol.Required(CONF_TO): deserialize_to_time,
    vol.Optional(CONF_DATA): CUSTOM_DATA_SCHEMA,
}
 
# Serialize time in validated config
STORAGE_TIME_RANGE_SCHEMA = vol.Schema(
    {
        vol.Required(CONF_FROM): vol.Coerce(str),
        vol.Required(CONF_TO): serialize_to_time,
        vol.Optional(CONF_DATA): CUSTOM_DATA_SCHEMA,
    }
)
 
SCHEDULE_SCHEMA: VolDictType = {
    vol.Optional(day, default=[]): vol.All(
        cv.ensure_list, [TIME_RANGE_SCHEMA], valid_schedule
    )
    for day in CONF_ALL_DAYS
}
STORAGE_SCHEDULE_SCHEMA: VolDictType = {
    vol.Optional(day, default=[]): vol.All(
        cv.ensure_list, [TIME_RANGE_SCHEMA], valid_schedule, [STORAGE_TIME_RANGE_SCHEMA]
    )
    for day in CONF_ALL_DAYS
}
 
# Validate YAML config
CONFIG_SCHEMA = vol.Schema(
    {DOMAIN: cv.schema_with_slug_keys(vol.All(BASE_SCHEMA | SCHEDULE_SCHEMA))},
    extra=vol.ALLOW_EXTRA,
)
# Validate storage config
STORAGE_SCHEMA = vol.Schema(
    {vol.Required(CONF_ID): cv.string} | BASE_SCHEMA | STORAGE_SCHEDULE_SCHEMA
)
# Validate + transform entity config
ENTITY_SCHEMA = vol.Schema(
    {vol.Required(CONF_ID): cv.string} | BASE_SCHEMA | SCHEDULE_SCHEMA
)
 
 
async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
    """Set up a schedule."""
    component = EntityComponent[Schedule](LOGGER, DOMAIN, hass)
 
    id_manager = IDManager()
 
    yaml_collection = YamlCollection(LOGGER, id_manager)
    sync_entity_lifecycle(hass, DOMAIN, DOMAIN, component, yaml_collection, Schedule)
 
    storage_collection = ScheduleStorageCollection(
        Store(
            hass,
            key=DOMAIN,
            version=STORAGE_VERSION,
            minor_version=STORAGE_VERSION_MINOR,
        ),
        id_manager,
    )
    sync_entity_lifecycle(hass, DOMAIN, DOMAIN, component, storage_collection, Schedule)
 
    await yaml_collection.async_load(
        [{CONF_ID: id_, **cfg} for id_, cfg in config.get(DOMAIN, {}).items()]
    )
    await storage_collection.async_load()
 
    DictStorageCollectionWebsocket(
        storage_collection,
        DOMAIN,
        DOMAIN,
        BASE_SCHEMA | STORAGE_SCHEDULE_SCHEMA,
        BASE_SCHEMA | STORAGE_SCHEDULE_SCHEMA,
    ).async_setup(hass)
 
    async def reload_service_handler(service_call: ServiceCall) -> None:
        """Reload yaml entities."""
        conf = await component.async_prepare_reload(skip_reset=True)
        if conf is None:
            conf = {DOMAIN: {}}
        await yaml_collection.async_load(
            [{CONF_ID: id_, **cfg} for id_, cfg in conf.get(DOMAIN, {}).items()]
        )
 
    async_register_admin_service(
        hass,
        DOMAIN,
        SERVICE_RELOAD,
        reload_service_handler,
    )
 
    return True
 
 
class ScheduleStorageCollection(DictStorageCollection):
    """Schedules stored in storage."""
 
    SCHEMA = vol.Schema(BASE_SCHEMA | STORAGE_SCHEDULE_SCHEMA)
 
    async def _process_create_data(self, data: dict) -> dict:
        """Validate the config is valid."""
        self.SCHEMASCHEMA(data)
        return data
 
    @callback
    def _get_suggested_id(self, info: dict) -> str:
        """Suggest an ID based on the config."""
        name: str = info[CONF_NAME]
        return name
 
    async def _update_data(self, item: dict, update_data: dict) -> dict:
        """Return a new updated data object."""
        self.SCHEMASCHEMA(update_data)
        return item | update_data
 
    async def _async_load_data(self) -> SerializedStorageCollection | None:
        """Load the data."""
        if data := await super()._async_load_data():
            data["items"] = [STORAGE_SCHEMA(item) for item in data["items"]]
        return data
 
 
class Schedule(CollectionEntity):
    """Schedule entity."""
 
    _entity_component_unrecorded_attributes = frozenset(
        {ATTR_EDITABLE, ATTR_NEXT_EVENT}
    )
 
    _attr_has_entity_name = True
    _attr_should_poll = False
    _attr_state: Literal["on", "off"]
    _config: ConfigType
    _next: datetime
    _unsub_update: Callable[[], None] | None = None
 
    def __init__(self, config: ConfigType, editable: bool) -> None:
        """Initialize a schedule."""
        self._config_config = ENTITY_SCHEMA(config)
        self._attr_capability_attributes_attr_capability_attributes = {ATTR_EDITABLE: editable}
        self._attr_icon_attr_icon = self._config_config.get(CONF_ICON)
        self._attr_name_attr_name = self._config_config[CONF_NAME]
        self._attr_unique_id_attr_unique_id = self._config_config[CONF_ID]
 
        # Exclude any custom attributes that may be present on time ranges from recording.
        self._unrecorded_attributes_unrecorded_attributes = self.all_custom_data_keysall_custom_data_keys()
        self._Entity__combined_unrecorded_attributes_Entity__combined_unrecorded_attributes = (
            self._entity_component_unrecorded_attributes_entity_component_unrecorded_attributes | self._unrecorded_attributes_unrecorded_attributes
        )
 
    @classmethod
    def from_storage(cls, config: ConfigType) -> Schedule:
        """Return entity instance initialized from storage."""
        return cls(config, editable=True)
 
    @classmethod
    def from_yaml(cls, config: ConfigType) -> Schedule:
        """Return entity instance initialized from yaml."""
        schedule = cls(config, editable=False)
        schedule.entity_id = f"{DOMAIN}.{config[CONF_ID]}"
        return schedule
 
    async def async_update_config(self, config: ConfigType) -> None:
        """Handle when the config is updated."""
        self._config_config = ENTITY_SCHEMA(config)
        self._attr_icon_attr_icon = config.get(CONF_ICON)
        self._attr_name_attr_name = config[CONF_NAME]
        self._clean_up_listener_clean_up_listener()
        self._update_update()
 
    @callback
    def _clean_up_listener(self) -> None:
        """Remove the update timer."""
        if self._unsub_update_unsub_update is not None:
            self._unsub_update_unsub_update()
            self._unsub_update_unsub_update = None
 
    async def async_added_to_hass(self) -> None:
        """Run when entity about to be added to hass."""
        self.async_on_removeasync_on_remove(self._clean_up_listener_clean_up_listener)
        self._update_update()
 
    @callback
    def _update(self, _: datetime | None = None) -> None:
        """Update the states of the schedule."""
        now = dt_util.now()
        todays_schedule = self._config_config.get(WEEKDAY_TO_CONF[now.weekday()], [])
 
        # Determine current schedule state
        for time_range in todays_schedule:
            # The current time should be greater or equal to CONF_FROM.
            if now.time() < time_range[CONF_FROM]:
                continue
            # The current time should be smaller (and not equal) to CONF_TO.
            # Note that any time in the day is treated as smaller than time.max.
            if now.time() < time_range[CONF_TO] or time_range[CONF_TO] == time.max:
                self._attr_state_attr_state = STATE_ON
                current_data = time_range.get(CONF_DATA)
                break
        else:
            self._attr_state_attr_state = STATE_OFF
            current_data = None
 
        # Find next event in the schedule, loop over each day (starting with
        # the current day) until the next event has been found.
        next_event = None
        for day in range(8):  # 8 because we need to search today's weekday next week
            day_schedule = self._config_config.get(
                WEEKDAY_TO_CONF[(now.weekday() + day) % 7], []
            )
            times = sorted(
                itertools.chain(
                    *[
                        [time_range[CONF_FROM], time_range[CONF_TO]]
                        for time_range in day_schedule
                    ]
                )
            )
 
            if next_event := next(
                (
                    possible_next_event
                    for timestamp in times
                    if (
                        possible_next_event := (
                            datetime.combine(now.date(), timestamp, tzinfo=now.tzinfo)
                            + timedelta(days=day)
                            if timestamp != time.max
                            # Special case for midnight of the following day.
                            else datetime.combine(now.date(), time(), tzinfo=now.tzinfo)
                            + timedelta(days=day + 1)
                        )
                    )
                    > now
                ),
                None,
            ):
                # We have found the next event in this day, stop searching.
                break
 
        self._attr_extra_state_attributes_attr_extra_state_attributes = {
            ATTR_NEXT_EVENT: next_event,
        }
 
        if current_data:
            # Add each key/value pair in the data to the entity's state attributes
            self._attr_extra_state_attributes_attr_extra_state_attributes.update(current_data)
 
        self.async_write_ha_stateasync_write_ha_state()
 
        if next_event:
            self._unsub_update_unsub_update = async_track_point_in_utc_time(
                self.hasshass,
                self._update_update,
                next_event,
            )
 
    def all_custom_data_keys(self) -> frozenset[str]:
        """Return the set of all currently used custom data attribute keys."""
        data_keys = set()
 
        for weekday in WEEKDAY_TO_CONF.values():
            if not (weekday_config := self._config_config.get(weekday)):
                continue  # this weekday is not configured
 
            for time_range in weekday_config:
                time_range_custom_data = time_range.get(CONF_DATA)
 
                if not time_range_custom_data or not isinstance(
                    time_range_custom_data, dict
                ):
                    continue  # this time range has no custom data, or it is not a dict
 
                data_keys.update(time_range_custom_data.keys())
 
        return frozenset(data_keys)