Poor Man's Use State Hook

poor_mans_use_state.py

from contextlib import contextmanager


def use_state(value):
    hook = current_hook()
    state = hook.use_state(lambda: {"value": value})
    return state["value"], lambda new_value: state.__setitem__("value", new_value)


def current_hook():
    assert _CURRENT_HOOK is not None, "No hook is currently active"
    return _CURRENT_HOOK


_CURRENT_HOOK = None


class Hook:
    def __init__(self):
        self._initialized = False
        self._state_index = 0
        self._state = []

    @contextmanager
    def render_context(self):
        global _CURRENT_HOOK
        self._state_index = 0
        _CURRENT_HOOK = self
        try:
            yield
        finally:
            self._initialized = True
            _CURRENT_HOOK = None

    def use_state(self, constructor):
        if not self._initialized:
            value = constructor()
            self._state.append(value)
        else:
            value = self._state[self._state_index]
        self._state_index += 1
        return value


# --- Demo ---


hook = Hook()

with hook.render_context():
    first, set_first = use_state(0)
    assert first == 0
    second, set_second = use_state(0)
    assert second == 0

with hook.render_context():
    first, set_first = use_state(0)
    assert first == 0
    second, set_second = use_state(0)
    assert second == 0

set_second(second + 1)

with hook.render_context():
    first, set_first = use_state(0)
    assert first == 0
    second, set_second = use_state(0)
    assert second == 1