supervisor/hassio/misc/scheduler.py

"""Schedule for Hass.io."""
import asyncio
from datetime import date, datetime, time, timedelta
import logging

_LOGGER: logging.Logger = logging.getLogger(__name__)

INTERVAL = "interval"
REPEAT = "repeat"
CALL = "callback"
TASK = "task"


class Scheduler:
    """Schedule task inside Hass.io."""

    def __init__(self):
        """Initialize task schedule."""
        self.loop = asyncio.get_running_loop()
        self._data = {}
        self.suspend = False

    def register_task(self, coro_callback, interval, repeat=True):
        """Schedule a coroutine.

        The coroutine need to be a callback without arguments.
        """
        task_id = hash(coro_callback)

        # Generate data
        opts = {CALL: coro_callback, INTERVAL: interval, REPEAT: repeat}

        # Schedule task
        self._data[task_id] = opts
        self._schedule_task(interval, task_id)

        return task_id

    def _run_task(self, task_id):
        """Run a scheduled task."""
        data = self._data[task_id]

        if not self.suspend:
            self.loop.create_task(data[CALL]())

        if data[REPEAT]:
            self._schedule_task(data[INTERVAL], task_id)
        else:
            self._data.pop(task_id)

    def _schedule_task(self, interval, task_id):
        """Schedule a task on loop."""
        if isinstance(interval, (int, float)):
            job = self.loop.call_later(interval, self._run_task, task_id)
        elif isinstance(interval, time):
            today = datetime.combine(date.today(), interval)
            tomorrow = datetime.combine(date.today() + timedelta(days=1), interval)

            # Check if we run it today or next day
            if today > datetime.today():
                calc = today
            else:
                calc = tomorrow

            job = self.loop.call_at(calc.timestamp(), self._run_task, task_id)
        else:
            _LOGGER.fatal(
                "Unknown interval %s (type: %s) for scheduler %s",
                interval,
                type(interval),
                task_id,
            )

        # Store job
        self._data[task_id][TASK] = job