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

.github/workflows/codequality.yaml

@@ -5,7 +5,7 @@ on:
 
 jobs:
   code-quality:
-    runs-on: ubuntu-20.04
+    runs-on: ubuntu-22.04
     strategy:
       fail-fast: false
       matrix:

.github/workflows/installation.yml

@@ -10,10 +10,10 @@ jobs:
 
     steps:
       - uses: actions/checkout@v4
-      - name: Set up Python 3.7
+      - name: Set up Python 3.8
         uses: actions/setup-python@v5
         with:
-          python-version: 3.7
+          python-version: 3.8
       - name: Install hatch
         run: pip install hatch
       - name: Build

.github/workflows/unittest.yml

@@ -17,11 +17,11 @@ jobs:
 
   unit-test:
     needs: [build, code-quality]
-    runs-on: ubuntu-20.04
+    runs-on: ubuntu-22.04
     strategy:
       fail-fast: false
       matrix:
-        python-version: [3.6, 3.7, 3.8, 3.9, "3.10", "3.11", "3.12", "3.13"]
+        python-version: [3.8, 3.9, "3.10", "3.11", "3.12", "3.13"]
 
     steps:
       - uses: actions/checkout@v4

pyproject.toml

@@ -16,8 +16,6 @@ dependencies = [
 
 classifiers = [
     "License :: OSI Approved :: MIT License",
-    "Programming Language :: Python :: 3.6",
-    "Programming Language :: Python :: 3.7",
     "Programming Language :: Python :: 3.8",
     "Programming Language :: Python :: 3.9",
     "Programming Language :: Python :: 3.10",
@@ -41,7 +39,7 @@ packages = ["react_ipywidgets", "reacton"]
 
 [project.optional-dependencies]
 dev = [
-    "ruff; python_version > '3.6'",
+    "ruff",
     "mypy",
     "pre-commit",
     "coverage",
@@ -57,7 +55,7 @@ dev = [
 ]
 
 generate = [
-    "ruff; python_version > '3.6'",
+    "ruff",
     "black",
     "bqplot",
     "jinja2",

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()