add tests for keyed hook preservation

this needs to be cleaned up now that we know what the bugs
were

Author: rmorshea

src/idom/core/component.py

@@ -21,8 +21,8 @@ def component(function: ComponentRenderFunction) -> Callable[..., "Component"]:
     """
 
     @wraps(function)
-    def constructor(*args: Any, **kwargs: Any) -> Component:
-        return Component(function, args, kwargs)
+    def constructor(*args: Any, key: str = "", **kwargs: Any) -> Component:
+        return Component(function, key, args, kwargs)
 
     return constructor
 
@@ -48,13 +48,16 @@ class Component(AbstractComponent):
     def __init__(
         self,
         function: ComponentRenderFunction,
+        key: str,
         args: Tuple[Any, ...],
         kwargs: Dict[str, Any],
     ) -> None:
-        self.key = kwargs.get("key", "")
+        self.key = key
         self._function = function
         self._args = args
         self._kwargs = kwargs
+        if key:
+            kwargs["key"] = key
 
     def render(self) -> VdomDict:
         model = self._function(*self._args, **self._kwargs)

src/idom/core/hooks.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import asyncio
+import weakref
 from logging import getLogger
 from threading import get_ident as get_thread_id
 from typing import (
@@ -374,7 +375,7 @@ class LifeCycleHook:
 
     __slots__ = (
         "component",
-        "_schedule_render_callback",
+        "_layout",
         "_schedule_render_later",
         "_current_state_index",
         "_state",
@@ -390,7 +391,7 @@ def __init__(
         layout: idom.core.layout.Layout,
     ) -> None:
         self.component = component
-        self._schedule_render_callback = layout.update
+        self._layout = weakref.ref(layout)
         self._schedule_render_later = False
         self._is_rendering = False
         self._rendered_atleast_once = False
@@ -468,4 +469,6 @@ def unset_current(self) -> None:
         del _current_life_cycle_hook[get_thread_id()]
 
     def _schedule_render(self) -> None:
-        self._schedule_render_callback(self.component)
+        layout = self._layout()
+        assert layout is not None
+        layout.update(self.component)

src/idom/core/layout.py

@@ -15,7 +15,7 @@
     Set,
     Tuple,
 )
-from weakref import ReferenceType, ref
+from weakref import ref
 
 from jsonpatch import apply_patch, make_patch
 from typing_extensions import TypedDict
@@ -125,7 +125,7 @@ async def _model_state_by_component_id(
 
     def _create_layout_update(self, component: AbstractComponent) -> LayoutUpdate:
         old_state = self._model_state_by_component_id[id(component)]
-        new_state = old_state.new()
+        new_state = old_state.new(None, component)
 
         self._render_component(old_state, new_state, component)
         changes = make_patch(getattr(old_state, "model", {}), new_state.model).patch
@@ -269,26 +269,33 @@ def _render_model_children(
 
         old_keys = set(old_state.children_by_key).difference(raw_typed_children_by_key)
         old_child_states = {key: old_state.children_by_key[key] for key in old_keys}
-        if old_child_states:
+        if old_keys:
             self._unmount_model_states(list(old_child_states.values()))
 
         new_children = new_state.model["children"] = []
         for index, (key, (child_type, child)) in enumerate(
             raw_typed_children_by_key.items()
         ):
             if child_type is DICT_TYPE:
-                child_state = _ModelState(ref(new_state), index, key, None)
-                self._render_model(old_child_states.get(key), child_state, child)
-                new_children.append(child_state.model)
-                new_state.children_by_key[key] = child_state
+                old_child_state = old_state.children_by_key.get(key)
+                if old_child_state is not None:
+                    new_child_state = old_child_state.new(new_state, None)
+                else:
+                    new_child_state = _ModelState(new_state, index, key, None)
+                self._render_model(old_child_state, new_child_state, child)
+                new_children.append(new_child_state.model)
+                new_state.children_by_key[key] = new_child_state
             elif child_type is COMPONENT_TYPE:
-                key = getattr(child, "key", "") or hex(id(child))
-                life_cycle_hook = LifeCycleHook(child, self)
-                child_state = _ModelState(ref(new_state), index, key, life_cycle_hook)
-                self._render_component(old_child_states.get(key), child_state, child)
-                new_children.append(child_state.model)
-                new_state.children_by_key[key] = child_state
-                self._model_state_by_component_id[id(child)] = child_state
+                old_child_state = old_state.children_by_key.get(key)
+                if old_child_state is not None:
+                    new_child_state = old_child_state.new(new_state, child)
+                else:
+                    hook = LifeCycleHook(child, self)
+                    new_child_state = _ModelState(new_state, index, key, hook)
+                self._render_component(old_child_state, new_child_state, child)
+                new_children.append(new_child_state.model)
+                new_state.children_by_key[key] = new_child_state
+                self._model_state_by_component_id[id(child)] = new_child_state
             else:
                 new_children.append(child)
 
@@ -299,14 +306,14 @@ def _render_model_children_without_old_state(
         for index, child in enumerate(raw_children):
             if isinstance(child, dict):
                 key = child.get("key") or hex(id(child))
-                child_state = _ModelState(ref(new_state), index, key, None)
+                child_state = _ModelState(new_state, index, key, None)
                 self._render_model(None, child_state, child)
                 new_children.append(child_state.model)
                 new_state.children_by_key[key] = child_state
             elif isinstance(child, AbstractComponent):
                 key = getattr(child, "key", "") or hex(id(child))
                 life_cycle_hook = LifeCycleHook(child, self)
-                child_state = _ModelState(ref(new_state), index, key, life_cycle_hook)
+                child_state = _ModelState(new_state, index, key, life_cycle_hook)
                 self._render_component(None, child_state, child)
                 new_children.append(child_state.model)
                 new_state.children_by_key[key] = child_state
@@ -322,6 +329,10 @@ def _unmount_model_states(self, old_states: List[_ModelState]) -> None:
                 hook = state.life_cycle_hook
                 hook.component_will_unmount()
                 del self._model_state_by_component_id[id(hook.component)]
+                import gc
+
+                print(state)
+                print(gc.get_referrers(hook))
             to_unmount.extend(state.children_by_key.values())
 
     def __repr__(self) -> str:
@@ -333,7 +344,7 @@ class _ModelState:
     __slots__ = (
         "index",
         "key",
-        "parent_ref",
+        "_parent_ref",
         "life_cycle_hook",
         "patch_path",
         "key_path",
@@ -347,40 +358,49 @@ class _ModelState:
 
     def __init__(
         self,
-        parent_ref: Optional[ReferenceType[_ModelState]],
+        parent: Optional[_ModelState],
         index: int,
         key: str,
         life_cycle_hook: Optional[LifeCycleHook],
     ) -> None:
         self.index = index
         self.key = key
 
-        if parent_ref is not None:
-            self.parent_ref = parent_ref
-            # temporarilly hydrate for use below
-            parent = parent_ref()
+        if parent is not None:
+            self._parent_ref = ref(parent)
+            self.key_path = f"{parent.key_path}/{key}"
+            self.patch_path = f"{parent.patch_path}/children/{index}"
         else:
-            parent = None
+            self.key_path = self.patch_path = ""
 
         if life_cycle_hook is not None:
             self.life_cycle_hook = life_cycle_hook
 
-        if parent is None:
-            self.key_path = self.patch_path = ""
-        else:
-            self.key_path = f"{parent.key_path}/{key}"
-            self.patch_path = f"{parent.patch_path}/children/{index}"
-
         self.event_targets: Set[str] = set()
         self.children_by_key: Dict[str, _ModelState] = {}
 
-    def new(self) -> _ModelState:
-        return _ModelState(
-            getattr(self, "parent_ref", None),
-            self.index,
-            self.key,
-            getattr(self, "life_cycle_hook", None),
-        )
+    @property
+    def parent(self) -> _ModelState:
+        # An AttributeError here is ok. It's synonymous
+        # with the existance of 'parent' attribute
+        p = self._parent_ref()
+        assert p is not None, "detached model state"
+        return p
+
+    def new(
+        self,
+        new_parent: Optional[_ModelState],
+        component: Optional[AbstractComponent],
+    ) -> _ModelState:
+        if new_parent is None:
+            new_parent = getattr(self, "parent", None)
+        if hasattr(self, "life_cycle_hook"):
+            assert component is not None
+            life_cycle_hook = self.life_cycle_hook
+            life_cycle_hook.component = component
+        else:
+            life_cycle_hook = None
+        return _ModelState(new_parent, self.index, self.key, life_cycle_hook)
 
     def iter_children(self, include_self: bool = True) -> Iterator[_ModelState]:
         to_yield = [self] if include_self else []
@@ -390,6 +410,28 @@ def iter_children(self, include_self: bool = True) -> Iterator[_ModelState]:
             to_yield.extend(node.children_by_key.values())
 
 
+class _ComponentQueue:
+
+    __slots__ = "_loop", "_queue", "_pending"
+
+    def __init__(self) -> None:
+        self._loop = asyncio.get_event_loop()
+        self._queue: "asyncio.Queue[AbstractComponent]" = asyncio.Queue()
+        self._pending: Set[int] = set()
+
+    def put(self, component: AbstractComponent) -> None:
+        component_id = id(component)
+        if component_id not in self._pending:
+            self._pending.add(component_id)
+            self._loop.call_soon_threadsafe(self._queue.put_nowait, component)
+        return None
+
+    async def get(self) -> AbstractComponent:
+        component = await self._queue.get()
+        self._pending.remove(id(component))
+        return component
+
+
 class _ModelEventTarget(TypedDict):
     target: str
     preventDefault: bool  # noqa
@@ -415,25 +457,3 @@ class _ModelVdomRequired(TypedDict, total=True):
 
 class _ModelVdom(_ModelVdomRequired, _ModelVdomOptional):
     """A VDOM dictionary model specifically for use with a :class:`Layout`"""
-
-
-class _ComponentQueue:
-
-    __slots__ = "_loop", "_queue", "_pending"
-
-    def __init__(self) -> None:
-        self._loop = asyncio.get_event_loop()
-        self._queue: "asyncio.Queue[AbstractComponent]" = asyncio.Queue()
-        self._pending: Set[int] = set()
-
-    def put(self, component: AbstractComponent) -> None:
-        component_id = id(component)
-        if component_id not in self._pending:
-            self._pending.add(component_id)
-            self._loop.call_soon_threadsafe(self._queue.put_nowait, component)
-        return None
-
-    async def get(self) -> AbstractComponent:
-        component = await self._queue.get()
-        self._pending.remove(id(component))
-        return component

tests/test_core/test_layout.py

@@ -384,7 +384,7 @@ def Inner():
         )
 
 
-async def test_keyed_components_get_garbage_collected():
+async def test_hooks_for_keyed_components_get_garbage_collected():
     pop_item = idom.Ref(None)
     garbage_collect_items = []
 
@@ -396,8 +396,7 @@ def Outer():
 
     @idom.component
     def Inner(key):
-        component = idom.hooks.current_hook().component
-        finalize(component, lambda: garbage_collect_items.append(key))
+        finalize(idom.hooks.current_hook(), lambda: garbage_collect_items.append(key))
         return idom.html.div(key)
 
     async with idom.Layout(Outer()) as layout:
@@ -428,3 +427,26 @@ def ComponentReturnsDuplicateKeys():
 
     with pytest.raises(ValueError, match="Duplicate keys in"):
         raise next(iter(caplog.records)).exc_info[1]
+
+
+async def test_keyed_components_preserve_hook_on_parent_update():
+    outer_hook = HookCatcher()
+    inner_hook = HookCatcher()
+
+    @idom.component
+    @outer_hook.capture
+    def Outer():
+        return Inner(key=1)
+
+    @idom.component
+    @inner_hook.capture
+    def Inner(key):
+        return idom.html.div(key)
+
+    async with idom.Layout(Outer()) as layout:
+        await layout.render()
+        old_inner_hook = inner_hook.current
+
+        outer_hook.current.schedule_render()
+        await layout.render()
+        assert old_inner_hook is inner_hook.current