fix: get_widget could fail with 'x not found in all known widgets'

This happened when we could re-render, and there was a function
component in the render tree that wouldn't need a re-execution.
It would then reconciliate with the old element data (also its arguments)
and the reference to the old element would map to get widget that
would be used in get_widget.
However, a use_effect/use_event would have a reference to a new element
and reacton would not be able to find the widget.
We now update the args and kwargs of the old element to the new
elements.

Author: maartenbreddels

reacton/core.py

@@ -1428,7 +1428,7 @@ def render(self, element: Element, container: widgets.Widget = None):
 
                 try:
                     self._shared_elements_next = set()
-                    self._render(self.element, "/", parent_key=ROOT_KEY)
+                    self._render(self.element, "/", parent_key=ROOT_KEY, old_to_new={})
                     self.first_render = False
                 except BaseException:
                     self._is_rendering = False
@@ -1463,7 +1463,7 @@ def format(reason: RerenderReason):
                             self.context.exceptions_children = []
                             self.context.exceptions_self = []
 
-                            self._render(self.element, "/", parent_key=ROOT_KEY)
+                            self._render(self.element, "/", parent_key=ROOT_KEY, old_to_new={})
                             logger.info("Render done: %r %r", self._rerender_needed, self._rerender_needed_reasons[-1])
                             assert self.context is self.context_root
                             render_counts += 1
@@ -1565,7 +1565,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: dict[Element, Element] = {}):
         if not isinstance(element, Element):
             raise TypeError(f"Expected element, not {element}")
         # for tracking stale data/elements when using get_widget
@@ -1729,12 +1729,43 @@ def _render(self, element: Element, default_key: str, parent_key: str):
                 else:
                     root_element = context.root_element_next or context.root_element
 
+                    # We don't rerender the component, but that mean that the root_element
+                    # refers to elements from a previous render pass. We need to update those so
+                    # that we have a working get_widget()
+                    # We first store a map of key->element for the invoked element of previous
+                    # render pass
+                    key_to_element: dict[str, Element] = {}
+
+                    def _store_key_to_element(el, key, parent_key):
+                        key_to_element[key] = el
+
+                    self._visit_children_values(el.kwargs, "/", "/", _store_key_to_element)
+                    self._visit_children_values(el.args, "/", "/", _store_key_to_element)
+                    # Next, we go over all children in the new element, and we then build up
+                    # A mapping of element->element from old to new
+                    # we will mutate, so make a copy
+                    old_to_new = old_to_new.copy()
+
+                    def _store_old_to_new(el, key, parent_key):
+                        old_to_new[el] = key_to_element[key]
+
+                    assert el_prev is not None
+                    self._visit_children_values(el_prev.kwargs, "/", "/", _store_old_to_new)
+                    self._visit_children_values(el_prev.args, "/", "/", _store_old_to_new)
+
+                    def map_old_to_new(el, key, parent_key):
+                        return old_to_new.get(el, el)
+
+                    assert root_element is not None
+                    root_element.kwargs = self._visit_children_values(root_element.kwargs, "/", "/", map_old_to_new)
+                    root_element.args = self._visit_children_values(root_element.args, "/", "/", map_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:

reacton/core_test.py

@@ -3204,3 +3204,87 @@ def Test():
         set_state(1)
     rc.render(w.HTML(value="recover").key("HTML"))
     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()