Keep keyed children when their wrapper container goes away

reacton/core.py

@@ -1306,29 +1306,30 @@ def __call__(self):
 
 class _RenderContext:
     context: Optional[ComponentContext] = None
-    # the element that an element of another type replaces, and its context, while
-    # reconciliation removes it (see _remove_replaced)
-    _replaced: Optional[Tuple[ComponentContext, Element]] = None
+    # the element that reconciliation removes because the new tree replaced it by an element
+    # of another type, or no longer has it, and its context (see _remove_outgoing)
+    _outgoing: Optional[Tuple[ComponentContext, Element]] = None
 
-    def _remove_replaced(self, el_prev: Element, key: str, parent_key: str):
-        """Remove the element at key, which an element of another type replaces."""
+    def _remove_outgoing(self, el: Element, key: str, parent_key: str):
+        """Remove the element at key, which the new tree replaced or no longer has."""
         assert self.context is not None
-        replaced = self._replaced
-        self._replaced = (self.context, el_prev)
+        outgoing = self._outgoing
+        self._outgoing = (self.context, el)
         try:
-            self._remove_element(el_prev, key, parent_key=parent_key)
+            self._remove_element(el, key, parent_key=parent_key)
         finally:
-            self._replaced = replaced
+            self._outgoing = outgoing
 
     def _keep_keyed_child(self, context: ComponentContext, el: Element, key: str) -> bool:
         # An explicit key is the same in any container of a component, so a child with an
-        # explicit key that the new tree still uses (it moved out of the replaced element, or
-        # stays under the new one) is not removed with the replaced element: reconciliation
-        # updates it where the new tree has it. Shared elements keep the old behavior: their
-        # bookkeeping (_shared_elements, _shared_widgets) needs the removal.
-        replaced = self._replaced
+        # explicit key that the new tree still uses (it moved out of the outgoing element, for
+        # instance a wrapper that went away, or stays under the element that replaces it) is
+        # not removed with the outgoing element: reconciliation updates it where the new tree
+        # has it. Shared elements keep the old behavior: their bookkeeping (_shared_elements,
+        # _shared_widgets) needs the removal.
+        outgoing = self._outgoing
         return (
-            replaced is not None and context is replaced[0] and el is not replaced[1] and el._key is not None and not el.is_shared and key in context.used_keys
+            outgoing is not None and context is outgoing[0] and el is not outgoing[1] and el._key is not None and not el.is_shared and key in context.used_keys
         )
 
     def __init__(self, element: Element, container: widgets.Widget = None, children_trait="children", handle_error: bool = True, initial_state=None):
@@ -2118,7 +2119,7 @@ def _reconsolidate(self, el: Element, default_key: str, parent_key: str):
         try:
             if isinstance(el.component, ComponentFunction):
                 if el_prev and isinstance(el_prev.component, ComponentWidget):
-                    self._remove_replaced(el_prev, key, parent_key=parent_key)
+                    self._remove_outgoing(el_prev, key, parent_key=parent_key)
                 new_parent_key = join_key(parent_key, key)
                 try:
                     # TODO: test suite passes when this block if commented out
@@ -2164,7 +2165,7 @@ def _reconsolidate(self, el: Element, default_key: str, parent_key: str):
                     if removed:
                         for key_remove in removed:
                             el_remove = elements[key_remove]
-                            self._remove_element(el_remove, key_remove, parent_key)
+                            self._remove_outgoing(el_remove, key_remove, parent_key)
                     for effect_index, effect in enumerate(child_context.effects):
                         if effect.next:
                             # if we have a next, it means that effect itself is executed
@@ -2295,7 +2296,7 @@ def reconsolidate_children():
                 else:
                     assert el_prev is not None, "widget_previous is not None, but el_prev is"
                     logger.debug("Replacing widget: %r → %r %r", el_prev, el, key)
-                    self._remove_replaced(el_prev, key, parent_key=parent_key)
+                    self._remove_outgoing(el_prev, key, parent_key=parent_key)
                     kwargs = reconsolidate_children()
                     widget = None
                     if not context.exceptions_children:
@@ -2372,7 +2373,7 @@ def reconsolidate_children():
             if extra:
                 for key in list(extra):
                     if key in self.context.elements:
-                        self._remove_element(self.context.elements[key], key, parent_key=parent_key)
+                        self._remove_outgoing(self.context.elements[key], key, parent_key=parent_key)
 
             # keeping this for debugging
             # logger.debug("Current:")
@@ -2796,7 +2797,7 @@ def _reconsolidate(self, el: Element, default_key: str, parent_key: str):
 
                 if el_prev and isinstance(el_prev.component, ComponentWidget):
                     # a widget element was replaced by a component element at this key
-                    self._remove_replaced(el_prev, key, parent_key=parent_key)
+                    self._remove_outgoing(el_prev, key, parent_key=parent_key)
                 new_parent_key = join_key(parent_key, key)
                 try:
                     if el.is_shared and (el.args or el.kwargs):
@@ -2832,7 +2833,7 @@ def _reconsolidate(self, el: Element, default_key: str, parent_key: str):
                         logger.info("elements to be removed: %r", stale_keys)
                         for stale_key in stale_keys:
                             if stale_key in child_context.elements:
-                                self._remove_element(child_context.elements[stale_key], stale_key, new_parent_key)
+                                self._remove_outgoing(child_context.elements[stale_key], stale_key, new_parent_key)
 
                     self._process_effects(child_context, context)
 
@@ -2904,7 +2905,7 @@ def _reconsolidate(self, el: Element, default_key: str, parent_key: str):
                 else:
                     assert el_prev is not None, "widget_previous is not None, but el_prev is"
                     # a different widget type at the same key: replace
-                    self._remove_replaced(el_prev, key, parent_key=parent_key)
+                    self._remove_outgoing(el_prev, key, parent_key=parent_key)
                     kwargs = self._visit_children_values(el.kwargs, key, parent_key, self._reconsolidate)
                     widget = None
                     if not context.exceptions_children:
@@ -3013,7 +3014,15 @@ def _cleanup_stale_effects_for_context(self, context: "ComponentContext", parent
                 if stale_key not in context.elements or stale_key in context.fast_stale_effects_cleaned:
                     continue
                 self.context = context
-                self._cleanup_stale_effects(context.elements[stale_key], stale_key, parent_key)
+                el = context.elements[stale_key]
+                # keyed children that the new tree still uses stay, with their effects
+                # (see _keep_keyed_child)
+                outgoing = self._outgoing
+                self._outgoing = (context, el)
+                try:
+                    self._cleanup_stale_effects(el, stale_key, parent_key)
+                finally:
+                    self._outgoing = outgoing
         finally:
             self.context = context_prev
 
@@ -3026,6 +3035,8 @@ def _cleanup_stale_effects(self, el: Element, default_key: str, parent_key: str)
         assert context is not None
         if key in context.fast_stale_effects_cleaned:
             return
+        if self._keep_keyed_child(context, el, key):
+            return
         if el.is_shared and (el in self._shared_elements_next or el not in self._shared_elements):
             return
         context.fast_stale_effects_cleaned.add(key)
@@ -3188,7 +3199,7 @@ def _remove_stale_root_elements(self, parent_key):
         stale_keys = sorted(set(self.context_root.elements) - self.context_root.used_keys)
         for stale_key in stale_keys:
             if stale_key in self.context_root.elements:
-                self._remove_element(self.context_root.elements[stale_key], stale_key, parent_key)
+                self._remove_outgoing(self.context_root.elements[stale_key], stale_key, parent_key)
 
 
 def _render_context_class():

reacton/core_test.py

@@ -4245,3 +4245,62 @@ def Test():
         assert rc.find(widgets.Button).widget.comm is not None
     assert len(rc._shared_widgets) == 1
     rc.close()
+
+
+@pytest.mark.parametrize("same_element", [True, False])
+def test_keyed_child_out_of_removed_wrapper(same_element):
+    # the wrapper around a child with an explicit key goes away (and comes back): its
+    # removal must not remove (or close) the child, which is still in the tree
+    set_wrapped = lambda x: None  # noqa
+
+    @react.component
+    def Child():
+        return w.Button(description="child")
+
+    child_once = Child().key("x")
+
+    @react.component
+    def Test():
+        nonlocal set_wrapped
+        wrapped, set_wrapped = react.use_state(True)
+        child = child_once if same_element else Child().key("x")
+        return w.VBox(children=[w.VBox(children=[child]) if wrapped else child])
+
+    root, rc = react.render_fixed(Test(), handle_error=False)
+    for wrapped in [False, True, False]:
+        set_wrapped(wrapped)
+        button = rc.find(widgets.Button).widget
+        assert button.comm is not None
+        assert (root.children[0] is button) != wrapped
+    rc.close()
+
+
+def test_keyed_child_out_of_removed_wrapper_keeps_its_effects():
+    # the kept child stays mounted: its effect must not be cleaned up when the wrapper goes
+    # away (the fast renderer cleans up stale subtrees early), only when the child leaves
+    log = []
+    set_wrapped = lambda x: None  # noqa
+
+    @react.component
+    def Child():
+        def effect():
+            log.append("run")
+            return lambda: log.append("cleanup")
+
+        react.use_effect(effect, [])
+        return w.Button(description="child")
+
+    @react.component
+    def Test():
+        nonlocal set_wrapped
+        wrapped, set_wrapped = react.use_state(True)
+        child = Child().key("x")
+        return w.VBox(children=[w.Label(value="sibling"), w.VBox(children=[child]) if wrapped else child])
+
+    box, rc = react.render(Test(), handle_error=False)
+    assert log == ["run"]
+    set_wrapped(False)
+    set_wrapped(True)
+    assert log == ["run"]
+    rc.close()
+    assert log == ["run", "cleanup"]

reacton/fuzz_test.py

@@ -42,6 +42,17 @@ def Leaf(id):
             return reacton.Fragment(children=[w.Button(description=f"f{id}.{i}") for i in range(value % 3)])
         return w.Button(description=f"leaf {id} {value}")
 
+    @react.component
+    def EffectLeaf(id):
+        # a leaf with an effect: however its containers change, the effect must run once
+        # while it stays in the tree, and be cleaned up once when it leaves
+        def effect():
+            log.append(("keyed effect", id))
+            return lambda: log.append(("keyed cleanup", id))
+
+        react.use_effect(effect, [])
+        return Leaf(id)
+
     @react.component
     def Wrapper(id):
         return Leaf(id * 7 + 1)
@@ -75,14 +86,29 @@ def Mover(id):
         # that moves between two sibling containers whose types flip
         state, set_state = react.use_state(0)
         registry[id] = set_state
-        child = react.use_memo(lambda: Leaf(id * 3 + 2).key(f"moved {id}"), [])
+        child = react.use_memo(lambda: EffectLeaf(id * 3 + 2).key(f"moved {id}"), [])
         kind = h(id, state) % 8
         First = w.HBox if kind & 2 else w.VBox
         Second = w.HBox if kind & 4 else w.VBox
         if kind & 1:
             return w.VBox(children=[First(children=[child]).key("first"), Second(children=[]).key("second")])
         return w.VBox(children=[First(children=[]).key("first"), Second(children=[child]).key("second")])
 
+    @react.component
+    def Unwrapper(id):
+        # a child with an explicit key inside 0, 1 or 2 wrapper containers: when a wrapper
+        # goes away (or comes back), the child stays in the tree at another level
+        state, set_state = react.use_state(0)
+        registry[id] = set_state
+        memo_child = react.use_memo(lambda: EffectLeaf(id * 3 + 2).key(f"unwrapped {id}"), [])
+        kind = h(id, state) % 12
+        child = memo_child if kind & 1 else EffectLeaf(id * 3 + 2).key(f"unwrapped {id}")
+        Wrapper = w.HBox if kind & 2 else w.VBox
+        element = child
+        for i in range((kind >> 2) % 3):
+            element = Wrapper(children=[element])
+        return w.VBox(children=[element])
+
     @react.component
     def Node(id, depth):
         state, set_state = react.use_state(0)
@@ -114,6 +140,8 @@ def cleanup():
                 child = Catcher(child_id)
             elif r < 0.8:
                 child = Mover(child_id)
+            elif r < 0.9:
+                child = Unwrapper(child_id)
             else:
                 child = Leaf(child_id)
             if rnd.random() < 0.3:
@@ -147,7 +175,7 @@ def _run_random_updates(fast: bool, seed: int, steps: int, batches: Optional[Lis
     record = batches is None
     batches = [] if batches is None else batches
     choices = random.Random(seed)
-    results = []
+    results: List[Any] = []
     with unittest.mock.patch.dict(os.environ, {"REACTON_FAST": "1" if fast else "0"}):
         widget, rc = react.render_fixed(Node(1, 0), handle_error=False)
         assert isinstance(rc, core._RenderContextFast) == fast
@@ -162,7 +190,10 @@ def _run_random_updates(fast: bool, seed: int, steps: int, batches: Optional[Lis
                     if id in registry:
                         registry[id](value)
             results.append((_widget_signature(rc.last_root_widget), sorted(registry), list(log)))
+        log.clear()
         rc.close()
+        # every effect that ran is cleaned up once, in the same order
+        results.append(("close", list(log)))
     return results, batches
 
 
@@ -179,4 +210,4 @@ def test_renderers_agree_on_random_updates(seed):
     finally:
         core.logger.setLevel(level)
     for step, (a, b) in enumerate(zip(expected, got)):
-        assert a == b, f"step {step}, batch {batches[step]}"
+        assert a == b, f"step {step}, batch {batches[step]}" if step < len(batches) else "close"