Track contexts in hooks as state (#787)

rmorshea · web-flow · commit 7bebc4d01a32 · 2022-07-14T16:46:51.000-07:00
* 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
diff --git a/docs/source/about/changelog.rst b/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
diff --git a/src/idom/backend/flask.py b/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(
diff --git a/src/idom/backend/sanic.py b/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(
diff --git a/src/idom/backend/starlette.py b/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(
diff --git a/src/idom/backend/tornado.py b/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)
diff --git a/src/idom/core/_f_back.py b/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
diff --git a/src/idom/core/_thread_local.py b/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
diff --git a/src/idom/core/hooks.py b/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:
diff --git a/src/idom/core/layout.py b/src/idom/core/layout.py
diff --git a/src/idom/core/vdom.py b/src/idom/core/vdom.py
diff --git a/tests/test_core/test_hooks.py b/tests/test_core/test_hooks.py
diff --git a/tests/test_core/test_layout.py b/tests/test_core/test_layout.py

-Original file line number
+Diff line change
 import abc
 import asyncio
 from collections import Counter
 +from contextlib import ExitStack
 from functools import wraps
 from logging import getLogger
 from typing import (
     def _create_layout_update(self, old_state: _ModelState) -> LayoutUpdate:
         new_state = _copy_component_model_state(old_state)
+-
         component = new_state.life_cycle_state.component
 -        self._render_component(old_state, new_state, component)
 -        # hook effects must run after the update is complete
 -        for model_state in _iter_model_state_children(new_state):
 -            if model_state.is_component_state:
 -                model_state.life_cycle_state.hook.affect_layout_did_render()
 +        with ExitStack() as exit_stack:
 +            self._render_component(exit_stack, old_state, new_state, component)
         old_model: Optional[VdomJson]
         try:
     def _render_component(
         self,
 +        exit_stack: ExitStack,
         old_state: Optional[_ModelState],
         new_state: _ModelState,
         component: ComponentType,
         else:
             life_cycle_hook = life_cycle_state.hook
             life_cycle_hook.affect_component_will_render(component)
 +            exit_stack.callback(life_cycle_hook.affect_layout_did_render)
 +            life_cycle_hook.set_current()
             try:
 -                life_cycle_hook.set_current()
 -                try:
 -                    raw_model = component.render()
 -                finally:
 -                    life_cycle_hook.unset_current()
 +                raw_model = component.render()
                 # wrap the model in a fragment (i.e. tagName="") to ensure components have
                 # a separate node in the model state tree. This could be removed if this
                 # components are given a node in the tree some other way
                 wrapper_model: VdomDict = {"tagName": ""}
                 if raw_model is not None:
                     wrapper_model["children"] = [raw_model]
 -                self._render_model(old_state, new_state, wrapper_model)
 +                self._render_model(exit_stack, old_state, new_state, wrapper_model)
             except Exception as error:
                 logger.exception(f"Failed to render {component}")
                 new_state.model.current = {
                     ),
+                }
             finally:
 +                life_cycle_hook.unset_current()
                 life_cycle_hook.affect_component_did_render()
         try:
     def _render_model(
         self,
 +        exit_stack: ExitStack,
         old_state: Optional[_ModelState],
         new_state: _ModelState,
         raw_model: Any,
         if "importSource" in raw_model:
             new_state.model.current["importSource"] = raw_model["importSource"]
         self._render_model_attributes(old_state, new_state, raw_model)
 -        self._render_model_children(old_state, new_state, raw_model.get("children", []))
 +        self._render_model_children(
 +            exit_stack, old_state, new_state, raw_model.get("children", [])
 +        )
     def _render_model_attributes(
         self,
     def _render_model_children(
         self,
 +        exit_stack: ExitStack,
         old_state: Optional[_ModelState],
         new_state: _ModelState,
         raw_children: Any,
         if old_state is None:
             if raw_children:
 -                self._render_model_children_without_old_state(new_state, raw_children)
 +                self._render_model_children_without_old_state(
 +                    exit_stack, new_state, raw_children
 +                )
             return None
         elif not raw_children:
             self._unmount_model_states(list(old_state.children_by_key.values()))
                             new_state,
                             index,
+                        )
 -                self._render_model(old_child_state, new_child_state, child)
 +                self._render_model(exit_stack, old_child_state, new_child_state, child)
                 new_children.append(new_child_state.model.current)
                 new_state.children_by_key[key] = new_child_state
             elif child_type is _COMPONENT_TYPE:
                         child,
                         self._rendering_queue.put,
+                    )
 -                self._render_component(old_child_state, new_child_state, child)
 +                self._render_component(
 +                    exit_stack, old_child_state, new_child_state, child
 +                )
             else:
                 old_child_state = old_state.children_by_key.get(key)
                 if old_child_state is not None:
                     self._unmount_model_states([old_child_state])
                 new_children.append(child)
     def _render_model_children_without_old_state(
 -        self, new_state: _ModelState, raw_children: List[Any]
 +        self,
 +        exit_stack: ExitStack,
 +        new_state: _ModelState,
 +        raw_children: List[Any],
     ) -> None:
         child_type_key_tuples = list(_process_child_type_and_key(raw_children))
         for index, (child, child_type, key) in enumerate(child_type_key_tuples):
             if child_type is _DICT_TYPE:
                 child_state = _make_element_model_state(new_state, index, key)
 -                self._render_model(None, child_state, child)
 +                self._render_model(exit_stack, None, child_state, child)
                 new_children.append(child_state.model.current)
                 new_state.children_by_key[key] = child_state
             elif child_type is _COMPONENT_TYPE:
                 child_state = _make_component_model_state(
                     new_state, index, key, child, self._rendering_queue.put
+                )
 -                self._render_component(None, child_state, child)
 +                self._render_component(exit_stack, None, child_state, child)
             else:
                 new_children.append(child)
         return False
 -def _iter_model_state_children(model_state: _ModelState) -> Iterator[_ModelState]:
 -    yield model_state
 -    for child in model_state.children_by_key.values():
 -        yield from _iter_model_state_children(child)
+-
+-
 def _new_root_model_state(
     component: ComponentType, schedule_render: Callable[[_LifeCycleStateId], None]
 ) -> _ModelState:
-Original file line number
+Diff line change
@@ @@ -1,6 +1,5 @@ @@
 from __future__ import annotations
 -import inspect
 import logging
 from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple, cast
     VdomJson,
+)
 +from ._f_back import f_module_name
++
 logger = logging.getLogger()
         "element represented by a :class:`VdomDict`."
+    )
 -    frame = inspect.currentframe()
 -    if frame is not None and frame.f_back is not None and frame.f_back is not None:
 -        module = frame.f_back.f_globals.get("__name__")  # module in outer frame
 -        if module is not None:
 -            qualname = module + "." + tag
 -            constructor.__module__ = module
 -            constructor.__qualname__ = qualname
 +    module_name = f_module_name(1)
 +    if module_name:
 +        constructor.__module__ = module_name
 +        constructor.__qualname__ = f"{module_name}.{tag}"
     return constructor
-Original file line number
+Diff line change
@@ @@ -1,5 +1,4 @@ @@
 import asyncio
 -import re
 import pytest
 def test_context_repr():
 -    Context = idom.create_context(None)
 -    assert re.match(r"Context\(.*\)", repr(Context()))
+-
 -    MyContext = idom.create_context(None, name="MyContext")
 -    assert re.match(r"MyContext\(.*\)", repr(MyContext()))
 +    sample_context = idom.create_context(None)
 +    assert repr(sample_context()) == f"ContextProvider({sample_context})"
 async def test_use_context_only_renders_for_value_change():
 async def test_neighboring_contexts_do_not_conflict():
 -    LeftContext = idom.create_context(None, name="Left")
 -    RightContext = idom.create_context(None, name="Right")
 +    LeftContext = idom.create_context(None)
 +    RightContext = idom.create_context(None)
     set_left = idom.Ref()
     set_right = idom.Ref()
         with assert_idom_did_not_log(r"SomeComponent\(.*?\) message is 'bye'"):
             await layout.render()
++
++
 +async def test_conditionally_rendered_components_can_use_context():
 +    set_state = idom.Ref()
 +    used_context_values = []
 +    some_context = idom.create_context(None)
++
 +    @idom.component
 +    def SomeComponent():
 +        state, set_state.current = idom.use_state(True)
 +        if state:
 +            return FirstCondition()
 +        else:
 +            return SecondCondition()
++
 +    @idom.component
 +    def FirstCondition():
 +        used_context_values.append(idom.use_context(some_context) + "-1")
++
 +    @idom.component
 +    def SecondCondition():
 +        used_context_values.append(idom.use_context(some_context) + "-2")
++
 +    async with idom.Layout(some_context(SomeComponent(), value="the-value")) as layout:
 +        await layout.render()
 +        assert used_context_values == ["the-value-1"]
 +        set_state.current(False)
 +        await layout.render()
 +        assert used_context_values == ["the-value-1", "the-value-2"]
-Original file line number
+Diff line change
     reset in any `Child()` components but there was a bug where that happened.
     """
 -    effect_calls_without_state = []
 +    effect_calls_without_state = set()
     set_child_key_num = StaticEventHandler()
     did_call_effect = asyncio.Event()
         async def record_if_state_is_reset():
             if state:
                 return
 -            effect_calls_without_state.append(child_key)
 +            effect_calls_without_state.add(child_key)
             set_state(1)
             did_call_effect.set()
     async with idom.Layout(Parent()) as layout:
         await layout.render()
         await did_call_effect.wait()
 -        assert effect_calls_without_state == ["some-key", "key-0"]
 +        assert effect_calls_without_state == {"some-key", "key-0"}
         did_call_effect.clear()
         for i in range(1, 5):
             await layout.deliver(LayoutEvent(set_child_key_num.target, []))
             await layout.render()
 -            assert effect_calls_without_state == ["some-key", "key-0"]
 +            assert effect_calls_without_state == {"some-key", "key-0"}
             did_call_effect.clear()