esphome/esphome/coroutine.py

"""
ESPHome's coroutine system.

The Problem: When running the code generationg, components can depend on variables being registered.
For example, an i2c-based sensor would need the i2c bus component to first be declared before the
codegen can emit code using that variable (or otherwise the C++ won't compile).

ESPHome's codegen system solves this by using coroutine-like methods. When a component depends on
a variable, it waits for it to be registered using `await cg.get_variable()`. If the variable
hasn't been registered yet, control will be yielded back to another component until the variable
is registered. This leads to a topological sort, solving the dependency problem.

Importantly, ESPHome only uses the coroutine *syntax*, no actual asyncio event loop is running in
the background. This is so that we can ensure the order of execution is constant for the same
YAML configuration, thus main.cpp only has to be recompiled if the configuration actually changes.

There are two syntaxes for ESPHome coroutines ("old style" vs "new style" coroutines).

"new style" - This is very much like coroutines you might be used to:

```py
async def my_coroutine(config):
    var = await cg.get_variable(config[CONF_ID])
    await some_other_coroutine(xyz)
    return var
```

new style coroutines are `async def` methods that use `await` to await the result of another coroutine,
and can return values using a `return` statement.

"old style" - This was a hack for when ESPHome still had to run on python 2, but is still compatible

```py
@coroutine
def my_coroutine(config):
    var = yield cg.get_variable(config[CONF_ID])
    yield some_other_coroutine(xyz)
    yield var
```

Here everything is combined in `yield` expressions. You await other coroutines using `yield` and
the last `yield` expression defines what is returned.
"""

import collections
import functools
import heapq
import inspect
import logging
import types
from typing import Any, Callable
from collections.abc import Awaitable, Generator, Iterator

_LOGGER = logging.getLogger(__name__)


def coroutine(func: Callable[..., Any]) -> Callable[..., Awaitable[Any]]:
    """Decorator to apply to methods to convert them to ESPHome coroutines."""
    if getattr(func, "_esphome_coroutine", False):
        # If func is already a coroutine, do not re-wrap it (performance)
        return func
    if inspect.isasyncgenfunction(func):
        # Trade-off: In ESPHome, there's not really a use-case for async generators.
        # and during the transition to new-style syntax it will happen that a `yield`
        # is not replaced properly, so don't accept async generators.
        raise ValueError(
            f"Async generator functions are not allowed. "
            f"Please check whether you've replaced all yields with awaits/returns. "
            f"See {func} in {func.__module__}"
        )
    if inspect.iscoroutinefunction(func):
        # A new-style async-def coroutine function, no conversion needed.
        return func

    if inspect.isgeneratorfunction(func):

        @functools.wraps(func)
        def coro(*args, **kwargs):
            gen = func(*args, **kwargs)
            ret = yield from _flatten_generator(gen)
            return ret

    else:
        # A "normal" function with no `yield` statements, convert to generator
        # that includes a yield just so it's also a generator function
        @functools.wraps(func)
        def coro(*args, **kwargs):
            res = func(*args, **kwargs)
            yield
            return res

    # Add coroutine internal python flag so that it can be awaited from new-style coroutines.
    coro = types.coroutine(coro)
    # pylint: disable=protected-access
    coro._esphome_coroutine = True
    return coro


def coroutine_with_priority(priority: float):
    """Decorator to apply to functions to convert them to ESPHome coroutines.

    :param priority: priority with which to schedule the coroutine, higher priorities run first.
    """

    def decorator(func):
        coro = coroutine(func)
        coro.priority = priority
        return coro

    return decorator


def _flatten_generator(gen: Generator[Any, Any, Any]):
    to_send = None
    while True:
        try:
            # Run until next yield expression
            val = gen.send(to_send)
        except StopIteration as e:
            # return statement or end of function

            # From py3.3, return with a value is allowed in generators,
            # and return value is transported in the value field of the exception.
            # If we find a value in the exception, use that as the return value,
            # otherwise use the value from the last yield statement ("old style")
            ret = to_send if e.value is None else e.value
            return ret

        if isinstance(val, collections.abc.Awaitable):
            # yielded object that is awaitable (like `yield some_new_style_method()`)
            # yield from __await__() like actual coroutines would.
            to_send = yield from val.__await__()
        elif inspect.isgenerator(val):
            # Old style, like `yield cg.get_variable()`
            to_send = yield from _flatten_generator(val)
        else:
            # Could be the last expression from this generator, record this as the return value
            to_send = val
            # perform a yield so that expressions like `while some_condition(): yield None`
            # do not run without yielding control back to the top
            yield


class FakeAwaitable:
    """Convert a generator to an awaitable object.

    Needed for internals of `cg.get_variable`. There we can't use @coroutine because
    native coroutines await from types.coroutine() directly without yielding back control to the top
    (likely as a performance enhancement).

    If we instead wrap the generator in this FakeAwaitable, control is yielded back to the top
    (reason unknown).
    """

    def __init__(self, gen: Generator[Any, Any, Any]) -> None:
        self._gen = gen

    def __await__(self):
        ret = yield from self._gen
        return ret


@functools.total_ordering
class _Task:
    def __init__(
        self,
        priority: float,
        id_number: int,
        iterator: Iterator[None],
        original_function: Any,
    ):
        self.priority = priority
        self.id_number = id_number
        self.iterator = iterator
        self.original_function = original_function

    def with_priority(self, priority: float) -> "_Task":
        return _Task(priority, self.id_number, self.iterator, self.original_function)

    @property
    def _cmp_tuple(self) -> tuple[float, int]:
        return (-self.priority, self.id_number)

    def __eq__(self, other):
        return self._cmp_tuple == other._cmp_tuple

    def __ne__(self, other):
        return not (self == other)

    def __lt__(self, other):
        return self._cmp_tuple < other._cmp_tuple


class FakeEventLoop:
    """Emulate an asyncio EventLoop to run some registered coroutine jobs in sequence."""

    def __init__(self):
        self._pending_tasks: list[_Task] = []
        self._task_counter = 0

    def add_job(self, func, *args, **kwargs):
        """Add a job to the task queue,

        Optionally retrieves priority from the function object, and schedules according to that.
        """
        if inspect.iscoroutine(func):
            raise ValueError("Can only add coroutine functions, not coroutine objects")
        if inspect.iscoroutinefunction(func):
            coro = func
            gen = coro(*args, **kwargs).__await__()
        else:
            coro = coroutine(func)
            gen = coro(*args, **kwargs)
        prio = getattr(coro, "priority", 0.0)
        task = _Task(prio, self._task_counter, gen, func)
        self._task_counter += 1
        heapq.heappush(self._pending_tasks, task)

    def flush_tasks(self):
        """Run until all tasks have been completed.

        :raises RuntimeError: if a deadlock is detected.
        """
        i = 0
        while self._pending_tasks:
            i += 1
            if i > 1000000:
                # Detect deadlock/circular dependency by measuring how many times tasks have been
                # executed. On the big tests/test1.yaml we only get to a fraction of this, so
                # this shouldn't be a problem.
                raise RuntimeError(
                    "Circular dependency detected! "
                    "Please run with -v option to see what functions failed to "
                    "complete."
                )

            task: _Task = heapq.heappop(self._pending_tasks)
            _LOGGER.debug(
                "Running %s in %s (num %s)",
                task.original_function.__qualname__,
                task.original_function.__module__,
                task.id_number,
            )

            try:
                next(task.iterator)
                # Decrease priority over time, so that if this task is blocked
                # due to a dependency others will clear the dependency
                # This could be improved with a less naive approach
                new_task = task.with_priority(task.priority - 1)
                heapq.heappush(self._pending_tasks, new_task)
            except StopIteration:
                _LOGGER.debug(" -> finished")