Track contexts in hooks as state (#787)

* fix use_context

components which are newly added to the layout after the initial
render do not have access to contexts. we solve this by tracking
life cycle hooks in stack and copying the context providers from
parent to child hooks.

we also change how contexts are implemented - instead of needing
a create_context function which returns a new Context class, we
just return a Context object that constructs a ContextProvider
component. the earlier implementation was too clever and did not
add anything for it.

* add test

* add changelog entry

* fix mypy

* fix tornado dev server

* remove unused func

* improve test

* remove unused method

* fix style issues

* get clever with functions

* fix tests

* update changelog

* get coverage

Author: rmorshea

docs/source/about/changelog.rst

@@ -23,13 +23,19 @@ more info, see the :ref:`Contributor Guide <Creating a Changelog Entry>`.
 Unreleased
 ----------
 
-**Added**
+**Fixed**
 
-- :pull:`123` - ``asgiref`` as a dependency
+- :issue:`789` - Conditionally rendered components cannot use contexts
 
 **Changed**
 
 - :pull:`123` - set default timeout on playwright page for testing
+- :pull:`787` - Track contexts in hooks as state
+- :pull:`787` - remove non-standard ``name`` argument from ``create_context``
+
+**Added**
+
+- :pull:`123` - ``asgiref`` as a dependency
 
 
 v0.39.0

src/idom/backend/flask.py

@@ -37,9 +37,7 @@
 
 logger = logging.getLogger(__name__)
 
-ConnectionContext: type[Context[Connection | None]] = create_context(
-    None, "ConnectionContext"
-)
+ConnectionContext: Context[Connection | None] = create_context(None)
 
 
 def configure(

src/idom/backend/sanic.py

@@ -32,9 +32,7 @@
 
 logger = logging.getLogger(__name__)
 
-ConnectionContext: type[Context[Connection | None]] = create_context(
-    None, "ConnectionContext"
-)
+ConnectionContext: Context[Connection | None] = create_context(None)
 
 
 def configure(

src/idom/backend/starlette.py

@@ -30,9 +30,7 @@
 
 logger = logging.getLogger(__name__)
 
-WebSocketContext: type[Context[WebSocket | None]] = create_context(
-    None, "WebSocketContext"
-)
+WebSocketContext: Context[WebSocket | None] = create_context(None)
 
 
 def configure(

src/idom/backend/tornado.py

@@ -26,9 +26,7 @@
 from .utils import CLIENT_BUILD_DIR, safe_client_build_dir_path
 
 
-ConnectionContext: type[Context[Connection | None]] = create_context(
-    None, "ConnectionContext"
-)
+ConnectionContext: Context[Connection | None] = create_context(None)
 
 
 def configure(
@@ -67,8 +65,7 @@ async def serve_development_app(
 ) -> None:
     enable_pretty_logging()
 
-    # setup up tornado to use asyncio
-    AsyncIOMainLoop().install()
+    AsyncIOMainLoop.current().install()
 
     server = HTTPServer(app)
     server.listen(port, host)

src/idom/core/_f_back.py

@@ -0,0 +1,23 @@
+from __future__ import annotations
+
+import inspect
+from types import FrameType
+
+
+def f_module_name(index: int = 0) -> str:
+    frame = f_back(index + 1)
+    if frame is None:
+        return ""  # pragma: no cover
+    name = frame.f_globals.get("__name__", "")
+    assert isinstance(name, str), "Expected module name to be a string"
+    return name
+
+
+def f_back(index: int = 0) -> FrameType | None:
+    frame = inspect.currentframe()
+    while frame is not None:
+        if index < 0:
+            return frame
+        frame = frame.f_back
+        index -= 1
+    return None  # pragma: no cover

src/idom/core/_thread_local.py

@@ -20,6 +20,3 @@ def get(self) -> _StateType:
         else:
             state = self._state[thread]
         return state
-
-    def set(self, state: _StateType) -> None:
-        self._state[current_thread()] = state

src/idom/core/hooks.py

@@ -8,7 +8,6 @@
     Any,
     Awaitable,
     Callable,
-    ClassVar,
     Dict,
     Generic,
     List,
@@ -239,108 +238,94 @@ def use_debug_value(
         logger.debug(f"{current_hook().component} {new}")
 
 
-def create_context(
-    default_value: _StateType, name: str | None = None
-) -> type[Context[_StateType]]:
+def create_context(default_value: _StateType) -> Context[_StateType]:
     """Return a new context type for use in :func:`use_context`"""
 
-    class _Context(Context[_StateType]):
-        _default_value = default_value
+    def context(
+        *children: Any,
+        value: _StateType = default_value,
+        key: Key | None = None,
+    ) -> ContextProvider[_StateType]:
+        return ContextProvider(
+            *children,
+            value=value,
+            key=key,
+            type=context,
+        )
+
+    context.__qualname__ = "context"
+
+    return context
 
-    _Context.__name__ = name or "Context"
 
-    return _Context
+class Context(Protocol[_StateType]):
+    """Returns a :class:`ContextProvider` component"""
 
+    def __call__(
+        self,
+        *children: Any,
+        value: _StateType = ...,
+        key: Key | None = ...,
+    ) -> ContextProvider[_StateType]:
+        ...
 
-def use_context(context_type: type[Context[_StateType]]) -> _StateType:
+
+def use_context(context: Context[_StateType]) -> _StateType:
     """Get the current value for the given context type.
 
     See the full :ref:`Use Context` docs for more information.
     """
-    # We have to use a Ref here since, if initially context_type._current is None, and
-    # then on a subsequent render it is present, we need to be able to dynamically adopt
-    # that newly present current context. When we update it though, we don't need to
-    # schedule a new render since we're already rending right now. Thus we can't do this
-    # with use_state() since we'd incur an extra render when calling set_state.
-    context_ref: Ref[Context[_StateType] | None] = use_ref(None)
-
-    if context_ref.current is None:
-        provided_context = context_type._current.get()
-        if provided_context is None:
-            # Cast required because of: https://github.com/python/mypy/issues/5144
-            return cast(_StateType, context_type._default_value)
-        context_ref.current = provided_context
-
-    # We need the hook now so that we can schedule an update when
     hook = current_hook()
+    provider = hook.get_context_provider(context)
+
+    if provider is None:
+        # force type checker to realize this is just a normal function
+        assert isinstance(context, FunctionType), f"{context} is not a Context"
+        # __kwdefault__ can be None if no kwarg only parameters exist
+        assert context.__kwdefaults__ is not None, f"{context} has no 'value' kwarg"
+        # lastly check that 'value' kwarg exists
+        assert "value" in context.__kwdefaults__, f"{context} has no 'value' kwarg"
+        # then we can safely access the context's default value
+        return cast(_StateType, context.__kwdefaults__["value"])
 
-    context = context_ref.current
+    subscribers = provider._subscribers
 
     @use_effect
     def subscribe_to_context_change() -> Callable[[], None]:
-        def set_context(new: Context[_StateType]) -> None:
-            # We don't need to check if `new is not context_ref.current` because we only
-            # trigger this callback when the value of a context, and thus the context
-            # itself changes. Therefore we can always schedule a render.
-            context_ref.current = new
-            hook.schedule_render()
-
-        context.subscribers.add(set_context)
-        return lambda: context.subscribers.remove(set_context)
-
-    return context.value
-
+        subscribers.add(hook)
+        return lambda: subscribers.remove(hook)
 
-_UNDEFINED: Any = object()
+    return provider._value
 
 
-class Context(Generic[_StateType]):
-
-    # This should be _StateType instead of Any, but it can't due to this limitation:
-    # https://github.com/python/mypy/issues/5144
-    _default_value: ClassVar[Any]
-
-    _current: ClassVar[ThreadLocal[Context[Any] | None]]
-
-    def __init_subclass__(cls) -> None:
-        # every context type tracks which of its instances are currently in use
-        cls._current = ThreadLocal(lambda: None)
-
+class ContextProvider(Generic[_StateType]):
     def __init__(
         self,
         *children: Any,
-        value: _StateType = _UNDEFINED,
-        key: Key | None = None,
+        value: _StateType,
+        key: Key | None,
+        type: Context[_StateType],
     ) -> None:
         self.children = children
-        self.value: _StateType = self._default_value if value is _UNDEFINED else value
         self.key = key
-        self.subscribers: set[Callable[[Context[_StateType]], None]] = set()
-        self.type = self.__class__
+        self.type = type
+        self._subscribers: set[LifeCycleHook] = set()
+        self._value = value
 
     def render(self) -> VdomDict:
-        current_ctx = self.__class__._current
-
-        prior_ctx = current_ctx.get()
-        current_ctx.set(self)
-
-        def reset_ctx() -> None:
-            current_ctx.set(prior_ctx)
-
-        current_hook().add_effect(COMPONENT_DID_RENDER_EFFECT, reset_ctx)
-
+        current_hook().set_context_provider(self)
         return vdom("", *self.children)
 
-    def should_render(self, new: Context[_StateType]) -> bool:
-        if self.value is not new.value:
-            new.subscribers.update(self.subscribers)
-            for set_context in self.subscribers:
-                set_context(new)
+    def should_render(self, new: ContextProvider[_StateType]) -> bool:
+        if self._value is not new._value:
+            for hook in self._subscribers:
+                hook.set_context_provider(new)
+                hook.schedule_render()
             return True
         return False
 
     def __repr__(self) -> str:
-        return f"{type(self).__name__}({id(self)})"
+        return f"{type(self).__name__}({self.type})"
 
 
 _ActionType = TypeVar("_ActionType")
@@ -558,14 +543,14 @@ def _try_to_infer_closure_values(
 
 def current_hook() -> LifeCycleHook:
     """Get the current :class:`LifeCycleHook`"""
-    hook = _current_hook.get()
-    if hook is None:
+    hook_stack = _hook_stack.get()
+    if not hook_stack:
         msg = "No life cycle hook is active. Are you rendering in a layout?"
         raise RuntimeError(msg)
-    return hook
+    return hook_stack[-1]
 
 
-_current_hook: ThreadLocal[LifeCycleHook | None] = ThreadLocal(lambda: None)
+_hook_stack: ThreadLocal[list[LifeCycleHook]] = ThreadLocal(list)
 
 
 EffectType = NewType("EffectType", str)
@@ -630,9 +615,8 @@ class LifeCycleHook:
 
             hook.affect_component_did_render()
 
-            # This should only be called after any child components yielded by
-            # component_instance.render() have also been rendered because effects of
-            # this type must run after the full set of changes have been resolved.
+            # This should only be called after the full set of changes associated with a
+            # given render have been completed.
             hook.affect_layout_did_render()
 
             # Typically an event occurs and a new render is scheduled, thus begining
@@ -650,6 +634,7 @@ class LifeCycleHook:
 
     __slots__ = (
         "__weakref__",
+        "_context_providers",
         "_current_state_index",
         "_event_effects",
         "_is_rendering",
@@ -666,6 +651,7 @@ def __init__(
         self,
         schedule_render: Callable[[], None],
     ) -> None:
+        self._context_providers: dict[Context[Any], ContextProvider[Any]] = {}
         self._schedule_render_callback = schedule_render
         self._schedule_render_later = False
         self._is_rendering = False
@@ -700,6 +686,14 @@ def add_effect(self, effect_type: EffectType, function: Callable[[], None]) -> N
         """Trigger a function on the occurance of the given effect type"""
         self._event_effects[effect_type].append(function)
 
+    def set_context_provider(self, provider: ContextProvider[Any]) -> None:
+        self._context_providers[provider.type] = provider
+
+    def get_context_provider(
+        self, context: Context[_StateType]
+    ) -> ContextProvider[_StateType] | None:
+        return self._context_providers.get(context)
+
     def affect_component_will_render(self, component: ComponentType) -> None:
         """The component is about to render"""
         self.component = component
@@ -753,13 +747,16 @@ def set_current(self) -> None:
         This method is called by a layout before entering the render method
         of this hook's associated component.
         """
-        _current_hook.set(self)
+        hook_stack = _hook_stack.get()
+        if hook_stack:
+            parent = hook_stack[-1]
+            self._context_providers.update(parent._context_providers)
+        hook_stack.append(self)
 
     def unset_current(self) -> None:
         """Unset this hook as the active hook in this thread"""
         # this assertion should never fail - primarilly useful for debug
-        assert _current_hook.get() is self
-        _current_hook.set(None)
+        assert _hook_stack.get().pop() is self
 
     def _schedule_render(self) -> None:
         try: