Simple reactivity in Python.

reactive.py

"""
Simple reactivity system for Python.
Not thread safe.
License: MIT
"""

from collections.abc import Callable
from dataclasses import dataclass
from typing import Generic, Protocol, TypeVar, final

_T = TypeVar("_T")


@dataclass
class Version:
    "Keep track of changes. Trigger for cache invalidation."

    version: int = 0

    def update(self) -> None:
        self.version += 1


class Reactive(Protocol):
    "Any reactive object: a `Var` or a `Computed`."

    _version: Version


@final
class _NeverCalled:
    """Sentinel value for when the `Computed` was never called and we don't
    have a value yet."""


# Keep track of what reactive variables are being accessed when computing a
# 'computed' variable.
_TRACKER_STACK: list[Callable[[Reactive], None]] = []


class Var(Generic[_T]):
    """
    Reactive variable.
    """

    def __init__(self, initial_value: _T) -> None:
        self._value = initial_value
        self._version = Version()

    def __call__(self) -> _T:
        "Retrieve value from this reactive variable."
        # Tell top-level trackers that this `Var` was accessed.
        if _TRACKER_STACK:
            _TRACKER_STACK[-1](self)

        return self._value

    def set(self, new_value: _T) -> None:
        "Set new value in this reactive variable."
        self._value = new_value
        self._version.update()


class Computed(Generic[_T]):
    """
    Computed: cache output value for as long as none of the referenced reactive
    variables change.

    A computed is evaluated lazily. Only when it's being called, we check
    whether it should be recomputed.
    """

    def __init__(self, func: Callable[[], _T]) -> None:
        self.func = func

        # Current evaluated value of this 'computed'.
        self._value: _T | _NeverCalled = _NeverCalled()

        # Version of this computed. This will be updated if this value changes.
        self._version = Version()

        # List of dependencies, together with their version when they were
        # observed for the last time.
        self._dependencies: list[tuple[Reactive, int]] = []

    def __call__(self) -> _T:
        "Get computed value."
        # Tell top-level trackers that this `Computed` was accessed.
        if _TRACKER_STACK:
            _TRACKER_STACK[-1](self)

        # Recompute, but only if we never computed before, or if any of our
        # dependencies changed.
        if isinstance(self._value, _NeverCalled) or self._any_dependency_changed():
            self._recompute()

        assert not isinstance(self._value, _NeverCalled)
        return self._value

    def _any_dependency_changed(self) -> bool:
        for dependency, previous_version in self._dependencies:
            if dependency._version.version != previous_version:
                return True
        return False

    def _recompute(self) -> None:
        # Clear list of dependencies. Add callback for keeping track of direct
        # dependencies.
        self._dependencies = []

        def add_dependency(reactive: Reactive) -> None:
            self._dependencies.append((reactive, reactive._version.version))

        _TRACKER_STACK.append(add_dependency)
        try:
            result = self.func()
            self._value = result
            self._version.update()
        finally:
            assert _TRACKER_STACK[-1] == add_dependency, "Tracker stack corruption."
            _TRACKER_STACK.pop()


def main() -> None:
    a = Var(4)

    @Computed
    def double() -> int:
        print("compute double of:", a())
        return a() * 2

    @Computed
    def result() -> str:
        return f"result={double()}"

    print(result())
    a.set(10)
    print(result())
    print(result())
    print(result())


if __name__ == "__main__":
    main()