Keep keyed children that move out of a replaced container

reacton/core.py

@@ -1306,6 +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
+
+    def _remove_replaced(self, el_prev: Element, key: str, parent_key: str):
+        """Remove the element at key, which an element of another type replaces."""
+        assert self.context is not None
+        replaced = self._replaced
+        self._replaced = (self.context, el_prev)
+        try:
+            self._remove_element(el_prev, key, parent_key=parent_key)
+        finally:
+            self._replaced = replaced
+
+    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
+        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
+        )
 
     def __init__(self, element: Element, container: widgets.Widget = None, children_trait="children", handle_error: bool = True, initial_state=None):
         self.element = element
@@ -2094,7 +2118,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_element(el_prev, default_key=key, parent_key=parent_key)
+                    self._remove_replaced(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
@@ -2271,7 +2295,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_element(el_prev, key, parent_key=parent_key)
+                    self._remove_replaced(el_prev, key, parent_key=parent_key)
                     kwargs = reconsolidate_children()
                     widget = None
                     if not context.exceptions_children:
@@ -2367,6 +2391,9 @@ def _remove_element(self, el: Element, default_key: str, parent_key):
         context = self.context
         logger.debug("Remove: (%s, %s) %r", parent_key, key, el)
 
+        if self._keep_keyed_child(context, el, key):
+            return
+
         if el.is_shared:
             if el not in self._shared_elements:
                 return
@@ -2769,7 +2796,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_element(el_prev, default_key=key, parent_key=parent_key)
+                    self._remove_replaced(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):
@@ -2877,7 +2904,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_element(el_prev, key, parent_key=parent_key)
+                    self._remove_replaced(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:
@@ -3041,6 +3068,9 @@ def _remove_element(self, el: Element, default_key: str, parent_key):
         context = self.context
         assert context is not None
 
+        if self._keep_keyed_child(context, el, key):
+            return
+
         if el.is_shared:
             if el not in self._shared_elements:
                 # another use of this element keeps it alive, or it was already removed

reacton/core_test.py

@@ -4187,3 +4187,61 @@ def App():
         "new run",
     ]
     rc.close()
+
+
+@pytest.mark.parametrize("from_first", [True, False])
+def test_keyed_child_moves_out_of_replaced_container(from_first):
+    # an element with an explicit key (the same object in every render) moves to a sibling
+    # container, while the container it leaves is replaced by one of another type: the
+    # removal of the old container must not remove (or close) the moved child
+    set_moved = lambda x: None  # noqa
+
+    @react.component
+    def Child():
+        return w.Button(description="child")
+
+    child = Child().key("x")
+
+    @react.component
+    def Test():
+        nonlocal set_moved
+        moved, set_moved = react.use_state(False)
+        source = [child] if not moved else []
+        target = [child] if moved else []
+        Source = w.HBox if moved else w.VBox
+        containers = [Source(children=source).key("source"), w.VBox(children=target).key("target")]
+        if not from_first:
+            containers = containers[::-1]
+        return w.VBox(children=containers)
+
+    box, rc = react.render(Test(), handle_error=False)
+    button = rc.find(widgets.Button).widget
+    set_moved(True)
+    buttons = rc.find(widgets.Button)
+    assert len(buttons) == 1
+    assert buttons.widget.comm is not None
+    assert rc.find(widgets.HBox).widget.children == ()
+    set_moved(False)
+    assert rc.find(widgets.Button).widget.comm is not None
+    assert button.comm is not None
+    rc.close()
+
+
+def test_shared_keyed_child_in_replaced_container():
+    # a shared element keeps the old behavior: it is removed with the replaced container
+    # and made again, and no old shared mappings stay behind
+    set_flip = lambda x: None  # noqa
+
+    @react.component
+    def Test():
+        nonlocal set_flip
+        flip, set_flip = react.use_state(False)
+        child = w.Button(description="shared").key("x").shared()
+        return w.VBox(children=[(w.HBox if flip else w.VBox)(children=[child])])
+
+    box, rc = react.render(Test(), handle_error=False)
+    for flip in [True, False, True]:
+        set_flip(flip)
+        assert rc.find(widgets.Button).widget.comm is not None
+    assert len(rc._shared_widgets) == 1
+    rc.close()

reacton/fuzz_test.py

@@ -69,6 +69,20 @@ def set_value(value):
             return w.Label(value=f"caught {exception}")
         return w.HBox(children=[Thrower(id * 3 + 2), Leaf(id * 3 + 1)])
 
+    @react.component
+    def Mover(id):
+        # an element made once (the same object in every render) with an explicit key,
+        # 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}"), [])
+        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 Node(id, depth):
         state, set_state = react.use_state(0)
@@ -98,6 +112,8 @@ def cleanup():
                 child = Wrapper(child_id)
             elif r < 0.7:
                 child = Catcher(child_id)
+            elif r < 0.8:
+                child = Mover(child_id)
             else:
                 child = Leaf(child_id)
             if rnd.random() < 0.3: