A lightweight Kotlin-style StateFlow for Python using RxPY, a reactive state container that always holds a current value and broadcasts updates to subscribers.

stateflow.py

import multiprocessing
from threading import Lock
from typing import Callable, Generic, TypeVar

from reactivex import Observable
from reactivex import operators as ops
from reactivex.scheduler import ThreadPoolScheduler
from reactivex.subject import BehaviorSubject

T = TypeVar("T")


class StateFlow(Generic[T]):
    """
    A lightweight Kotlin-like StateFlow built on RxPY.

    Concept
    -------
    Think of this as a *reactive variable*.

    It behaves like:
        value = 10

    But observers can subscribe and react whenever the value changes.

    Core Guarantees
    ---------------
    1. A current value always exists.
    2. New subscribers instantly receive the latest value.
    3. Duplicate updates are ignored (same value won't re-emit).
    4. Observers run on a background thread pool so producers are not blocked.

    Internals
    ---------
    BehaviorSubject
        Holds the latest value and emits it to new subscribers.

    distinct_until_changed()
        Prevents emitting identical values.

    ThreadPoolScheduler
        Runs observers asynchronously so slow consumers do not block producers.

    Comparison with Kotlin
    ----------------------
    Kotlin:
        val state = MutableStateFlow(0)
        state.value = 10
        state.collect { println(it) }

    Python:
        state = StateFlow(0)
        state.set(10)
        state.subscribe(print)
    """

    def __init__(self, initial: T):
        # BehaviorSubject stores the latest value and immediately emits it
        # to new subscribers.
        self._subject = BehaviorSubject(initial)

        # Lock ensures atomic updates when using update()
        self._lock = Lock()

        # Use half of available CPUs for observer threads
        workers = max(1, multiprocessing.cpu_count() // 2)
        self._scheduler = ThreadPoolScheduler(workers)

        # Main observable stream
        # distinct_until_changed prevents duplicate emissions
        # observe_on ensures observers run asynchronously
        self._stream: Observable[T] = self._subject.pipe(
            ops.distinct_until_changed(),
            ops.observe_on(self._scheduler),
        )

    @property
    def value(self) -> T:
        """
        Returns the current state value.

        Equivalent to Kotlin:
            state.value
        """
        return self._subject.value

    def set(self, value: T) -> None:
        """
        Set a new value and emit it to subscribers.

        Example
        -------
        state.set(10)
        """
        self._subject.on_next(value)

    def update(self, fn: Callable[[T], T]) -> None:
        """
        Atomically update the state using the previous value.

        This avoids race conditions when multiple threads update state.

        Example
        -------
        state.update(lambda v: v + 1)
        """
        with self._lock:
            new_value = fn(self._subject.value)
            self._subject.on_next(new_value)

    def subscribe(self, observer):
        """
        Subscribe to state changes.

        The observer receives the current value immediately
        and every subsequent update.

        Example
        -------
        state.subscribe(lambda v: print("state:", v))
        """
        return self._stream.subscribe(observer)

    def pipe(self, *operators) -> Observable:
        """
        Apply Rx operators to the stream.

        This allows reactive transformations similar to Kotlin Flow.

        Example
        -------
        doubled = state.pipe(
            ops.map(lambda v: v * 2)
        )

        doubled.subscribe(print)
        """
        return self._stream.pipe(*operators)