Fix get_widget for a child element of a memoized component

reacton/core.py

@@ -1784,7 +1784,7 @@ def format(reason: RerenderReason):
                 raise exc
         return widget
 
-    def _render(self, element: Element, default_key: str, parent_key: str):
+    def _render(self, element: Element, default_key: str, parent_key: str, old_to_new: Optional[Dict[Element, Element]] = None):
         if not isinstance(element, Element):
             raise TypeError(f"Expected element, not {element}")
         # for tracking stale data/elements when using get_widget
@@ -1837,7 +1837,11 @@ def _render(self, element: Element, default_key: str, parent_key: str):
             logger.debug("Render: arguments... (children of %s,%s)", parent_key, key)
             # only when we landed at a widget leaf, or a shared element, we need to render the children
             if isinstance(el.component, ComponentWidget) or el.is_shared:
-                self._visit_children(el, key, parent_key, self._render)
+                if old_to_new:
+                    self._substitute_old_elements(el, old_to_new)
+                    self._visit_children(el, key, parent_key, functools.partial(self._render, old_to_new=old_to_new))
+                else:
+                    self._visit_children(el, key, parent_key, self._render)
             assert self.context is context
             logger.debug("Render: arguments done (children of %s,%s)", parent_key, key)
 
@@ -1913,6 +1917,9 @@ def _render(self, element: Element, default_key: str, parent_key: str):
                     # we reset if before calling the component
                     # which might set it to true again
                     context.needs_render = False
+                    # the component function runs with the arguments of the new invoke
+                    # element, so everything it returns is fresh: nothing left to substitute
+                    old_to_new = None
                     # Now, we actually execute the render function, and get
                     # back the root element
                     root_element: Optional[Element] = None
@@ -1953,12 +1960,16 @@ def _render(self, element: Element, default_key: str, parent_key: str):
                 else:
                     root_element = context.root_element_next or context.root_element
 
+                    assert el_prev is not None
+                    assert root_element is not None
+                    old_to_new = self._refresh_reused_root_element(el, el_prev, root_element, old_to_new)
+
                 if self.render_count != render_count:
                     raise RuntimeError("Recursive render detected, possible a bug in react")
                 if root_element is not None:
                     logger.debug("root element: %r %x", root_element, id(root_element))
                     new_parent_key = join_key(parent_key, key)
-                    self._render(root_element, "/", parent_key=new_parent_key)  # depth first
+                    self._render(root_element, "/", parent_key=new_parent_key, old_to_new=old_to_new)  # depth first
                     context.root_element_next = root_element
                 else:
                     if el.is_shared:
@@ -2398,6 +2409,57 @@ def _visit_children(self, el: Element, default_key: str, parent_key: str, f: Cal
         self._visit_children_values(el.kwargs, key, parent_key, f)
         self._visit_children_values(el.args, key, parent_key, f)
 
+    def _substitute_old_elements(self, el: Element, old_to_new: Dict[Element, Element]) -> None:
+        # el sits inside a reused root element, so its arguments can still refer to child
+        # elements of a previous render pass. _refresh_reused_root_element does not reach
+        # them, because the traversal stops at every element boundary.
+        def map_old_to_new(child: Element, key: str, parent_key: str):
+            return old_to_new.get(child, child)
+
+        el.kwargs = self._visit_children_values(el.kwargs, "/", "/", map_old_to_new)
+        el.args = self._visit_children_values(el.args, "/", "/", map_old_to_new)
+
+    def _refresh_reused_root_element(
+        self,
+        el: Element,
+        el_prev: Element,
+        root_element: Element,
+        old_to_new: Optional[Dict[Element, Element]],
+    ) -> Dict[Element, Element]:
+        # We reuse the root element of a component we did not re-execute. That root element
+        # refers to child elements of a previous render pass, while the caller (which did
+        # re-execute) holds new, equal-by-value element objects. Without a fixup, the widget
+        # ends up registered under the old element and get_widget(new element) fails.
+        # We build a map from old to new child element, and substitute them in the reused
+        # root element. The map is passed down, because a component deeper in the tree can
+        # reuse its root element as well.
+        substitutions: Dict[Element, Element] = {}
+        if el is not el_prev:
+            key_to_element: Dict[str, Element] = {}
+
+            def store_key_to_element(child: Element, key: str, parent_key: str):
+                key_to_element[key] = child
+
+            self._visit_children_values(el.kwargs, "/", "/", store_key_to_element)
+            self._visit_children_values(el.args, "/", "/", store_key_to_element)
+
+            if key_to_element:
+                # the arguments compare equal by value, so the traversal keys of the
+                # previous invoke element match those of the new one
+                def store_old_to_new(child: Element, key: str, parent_key: str):
+                    replacement = key_to_element[key]
+                    if replacement is not child:
+                        substitutions[child] = replacement
+
+                self._visit_children_values(el_prev.kwargs, "/", "/", store_old_to_new)
+                self._visit_children_values(el_prev.args, "/", "/", store_old_to_new)
+
+        if old_to_new:
+            substitutions = {**old_to_new, **substitutions}
+        if substitutions:
+            self._substitute_old_elements(root_element, substitutions)
+        return substitutions
+
     def _visit_children_values(self, value: Any, key: str, parent_key: str, f: Callable):
         if isinstance(value, Element):
             return f(value, key, parent_key)
@@ -2450,7 +2512,7 @@ def _set_rerender_needed(self, reason: str):
         self._rerender_needed_reasons.append(RerenderReason(reason=reason))
         self._rerender_needed = True
 
-    def _render(self, element: Element, default_key: str, parent_key: str):
+    def _render(self, element: Element, default_key: str, parent_key: str, old_to_new: Optional[Dict[Element, Element]] = None):
         if not isinstance(element, Element):
             raise TypeError(f"Expected element, not {element}")
         # for tracking stale elements when using get_widget
@@ -2494,22 +2556,31 @@ def _render(self, element: Element, default_key: str, parent_key: str):
                 del context.children_next[key]
             # the element arguments are part of this component's element tree
             if el.kwargs:
-                self._visit_children(el, key, parent_key, self._render)
+                if old_to_new:
+                    self._substitute_old_elements(el, old_to_new)
+                    self._visit_children(el, key, parent_key, functools.partial(self._render, old_to_new=old_to_new))
+                else:
+                    self._visit_children(el, key, parent_key, self._render)
             return
 
         assert isinstance(el.component, ComponentFunction)
         if el.is_shared and (el.args or el.kwargs):
             # arguments of a shared element belong to the context it is rendered in;
             # for non-shared component elements the component function decides
             # what ends up in the tree
-            self._visit_children(el, key, parent_key, self._render)
+            if old_to_new:
+                self._substitute_old_elements(el, old_to_new)
+                self._visit_children(el, key, parent_key, functools.partial(self._render, old_to_new=old_to_new))
+            else:
+                self._visit_children(el, key, parent_key, self._render)
 
         context_previous = context.children_next.get(key)
         if context_previous is None:
             context_previous = context.children.get(key)
 
         if (
             not self._walk_all
+            and not old_to_new
             and el is el_prev
             and not el.is_shared
             and context_previous is not None
@@ -2578,6 +2649,8 @@ def _render(self, element: Element, default_key: str, parent_key: str):
                 # we reset it before calling the component function,
                 # which might set it to true again
                 context.needs_render = False
+                # everything the component function returns is fresh, see the classic renderer
+                old_to_new = None
                 try:
                     with contextlib.ExitStack() as stack:
                         for cm in context.context_managers:
@@ -2612,14 +2685,17 @@ def _render(self, element: Element, default_key: str, parent_key: str):
                     raise ValueError(f"Component {el.component} returned None")
             else:
                 root_element = context.root_element_next or context.root_element
+                assert el_prev is not None
+                assert root_element is not None
+                old_to_new = self._refresh_reused_root_element(el, el_prev, root_element, old_to_new)
 
             if self.render_count != render_count_check:
                 raise RuntimeError("Recursive render detected, possible a bug in react")
 
             # the subtree walk below will mark this again when state changes
             context.needs_render_descendant = False
             if root_element is not None:
-                self._render(root_element, "/", parent_key=join_key(parent_key, key))  # depth first
+                self._render(root_element, "/", parent_key=join_key(parent_key, key), old_to_new=old_to_new)  # depth first
                 context.root_element_next = root_element
             elif el.is_shared:
                 self._shared_elements_next.remove(el)

reacton/core_test.py

@@ -3447,3 +3447,116 @@ def effect():
     set_other(1)
     assert root.children[0].description == "other 1"
     rc.close()
+
+
+def test_get_widget_fail_on_rerender_use_event():
+    @reacton.component
+    def Test():
+        force_rerender, set_force_rerender = react.use_state(0, key="force_rerender")
+        click_works, set_click_works = react.use_state(False, key="click_works")
+
+        with ContainerFunction():
+            el = v.Btn(children=["Works"] if click_works else ["Does not work"])
+
+        v.use_event(el, "click", lambda *_ignore: set_click_works(True))
+
+        if force_rerender == 0:
+            set_force_rerender(1)
+
+    box, rc = react.render(Test(), handle_error=False)
+    rc.find(ipyvuetify.Btn).widget.click()
+    assert rc.find(ipyvuetify.Btn).widget.children[0] == "Works"
+    rc.close()
+
+
+@pytest.mark.parametrize("on_use_effect", [True, False])
+def test_get_widget_fail_on_rerender_simple(on_use_effect):
+    @reacton.component
+    def Test():
+        force_rerender, set_force_rerender = react.use_state(0, key="force_rerender")
+
+        def effect():
+            widget = react.get_widget(el)
+
+            assert widget is not None
+
+        use_effect(effect, [])
+        with ContainerFunction():
+            el = w.Button(description="Hi")
+
+        react.use_effect(effect, None)
+
+        if force_rerender == 0 and not on_use_effect:
+            set_force_rerender(1)
+
+        def possibly_rerender():
+            if force_rerender == 0 and on_use_effect:
+                set_force_rerender(1)
+
+        use_effect(possibly_rerender, None)
+
+    box, rc = react.render(Test(), handle_error=False)
+    rc.close()
+
+
+def test_get_widget_fail_on_rerender_complex(Container1, Container2):
+    @reacton.component
+    def MakeMoreComplex(arg, children=[]):
+        return Container1(children=[*children, arg])
+
+    @reacton.component
+    def Test():
+        force_rerender, set_force_rerender = react.use_state(0, key="force_rerender")
+
+        def effect():
+            widget1 = react.get_widget(el1)
+            widget2 = react.get_widget(el2)
+
+            assert widget1 is not None
+            assert widget2 is not None
+
+        use_effect(effect, [])
+        el1 = w.Button(description="Foo")
+        with MakeMoreComplex(el1):
+            with Container2():
+                el2 = w.Button(description="Bar")
+
+        react.use_effect(effect, None)
+
+        def rerender():
+            if force_rerender == 0:
+                set_force_rerender(1)
+
+        react.use_effect(rerender, [])
+
+    box, rc = react.render(Test(), handle_error=False)
+    rc.close()
+
+
+def test_get_widget_fail_on_rerender_nested_widgets():
+    @react.component
+    def NestDeep(arg, children=[]):
+        # the received element sits two element levels below the root element,
+        # so a substitution that stops at the first element boundary misses it
+        return w.VBox(children=[w.HBox(children=[arg])])
+
+    @reacton.component
+    def Test():
+        force_rerender, set_force_rerender = react.use_state(0, key="force_rerender")
+
+        def effect():
+            assert react.get_widget(el) is not None
+
+        el = w.Button(description="Deep")
+        NestDeep(el)
+
+        react.use_effect(effect, None)
+
+        def rerender():
+            if force_rerender == 0:
+                set_force_rerender(1)
+
+        use_effect(rerender, [])
+
+    box, rc = react.render(Test(), handle_error=False)
+    rc.close()