Track contexts in hooks as state

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: