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.

Author: rmorshea

src/idom/__init__.py

@@ -4,6 +4,7 @@
 from .core.component import component
 from .core.events import event
 from .core.hooks import (
+    Context,
     create_context,
     use_callback,
     use_context,
@@ -27,6 +28,7 @@
 __all__ = [
     "component",
     "config",
+    "Context",
     "create_context",
     "event",
     "hooks",

src/idom/backend/flask.py

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

src/idom/backend/sanic.py

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

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, "WebSocketContext")
 
 
 def configure(

src/idom/backend/tornado.py

@@ -26,7 +26,7 @@
 from .utils import CLIENT_BUILD_DIR, safe_client_build_dir_path
 
 
-ConnectionContext: type[Context[Connection | None]] = create_context(
+ConnectionContext: Context[Connection | None] = create_context(
     None, "ConnectionContext"
 )
 

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 ""
+    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

src/idom/core/hooks.py

@@ -8,7 +8,6 @@
     Any,
     Awaitable,
     Callable,
-    ClassVar,
     Dict,
     Generic,
     List,
@@ -21,12 +20,14 @@
     cast,
     overload,
 )
+from warnings import warn
 
 from typing_extensions import Protocol
 
 from idom.config import IDOM_DEBUG_MODE
 from idom.utils import Ref
 
+from ._f_back import f_module_name
 from ._thread_local import ThreadLocal
 from .types import ComponentType, Key, VdomDict
 from .vdom import vdom
@@ -240,107 +241,94 @@ def use_debug_value(
 
 
 def create_context(
-    default_value: _StateType, name: str | None = None
-) -> type[Context[_StateType]]:
+    default_value: _StateType, name: str = "Context"
+) -> Context[_StateType]:
     """Return a new context type for use in :func:`use_context`"""
+    warn(
+        "The 'create_context' function is deprecated. "
+        "Use the 'Context' class instead. "
+        "This function will be removed in a future release."
+    )
+    return Context(name, default_value)
+
 
-    class _Context(Context[_StateType]):
-        _default_value = default_value
+_UNDEFINED = object()
 
-    _Context.__name__ = name or "Context"
 
-    return _Context
+class Context(Generic[_StateType]):
+    """Returns a :class:`ContextProvider` component"""
 
+    def __init__(self, name: str, default_value: _StateType) -> None:
+        self.name = name
+        self.default_value = default_value
 
-def use_context(context_type: type[Context[_StateType]]) -> _StateType:
+    def __call__(
+        self,
+        *children: Any,
+        value: _StateType = _UNDEFINED,
+        key: Key | None = None,
+    ) -> (
+        # users don't need to see that this is a ContextProvider
+        ComponentType
+    ):
+        return ContextProvider(
+            *children,
+            value=self.default_value if value is _UNDEFINED else value,
+            key=key,
+            type=self,
+        )
+
+    def __repr__(self):
+        return f"{type(self).__name__}({self.name!r})"
+
+
+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()
 
-    context = context_ref.current
+    provider = hook.get_context_provider(context)
+    if provider is None:
+        return context.default_value
 
     @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
-
+        provider.subscribers.add(hook)
+        return lambda: provider.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.value = value
         self.key = key
-        self.subscribers: set[Callable[[Context[_StateType]], None]] = set()
-        self.type = self.__class__
+        self.subscribers: set[LifeCycleHook] = set()
+        self.type = type
 
     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:
+    def should_render(self, new: ContextProvider[_StateType]) -> bool:
         if self.value is not new.value:
-            new.subscribers.update(self.subscribers)
-            for set_context in self.subscribers:
-                set_context(new)
+            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.name!r})"
 
 
 _ActionType = TypeVar("_ActionType")
@@ -558,14 +546,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 +618,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 +637,7 @@ class LifeCycleHook:
 
     __slots__ = (
         "__weakref__",
+        "_context_providers",
         "_current_state_index",
         "_event_effects",
         "_is_rendering",
@@ -666,6 +654,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 +689,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 +750,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: