improvement: hold on to unnamed UI elements

marimo/_plugins/ui/_core/registry.py

@@ -2,19 +2,20 @@
 from __future__ import annotations
 
 import sys
-import weakref
-from typing import Any, Dict, Iterable, Mapping, TypeVar, Union
+from typing import TYPE_CHECKING, Any, Dict, Iterable, Mapping, TypeVar, Union
 
 if sys.version_info < (3, 10):
     from typing_extensions import TypeAlias
 else:
     from typing import TypeAlias
 
 from marimo._ast.app import _Namespace
-from marimo._ast.cell import CellId_t
 from marimo._plugins.ui._core.ui_element import UIElement
 from marimo._runtime.context import get_context
 
+if TYPE_CHECKING:
+    from marimo._ast.cell import CellId_t
+
 UIElementId = str
 
 T = TypeVar("T")
@@ -26,11 +27,15 @@
 class UIElementRegistry:
     def __init__(self) -> None:
         # mapping from object id to UIElement object that has that id
-        self._objects: dict[UIElementId, weakref.ref[UIElement[Any, Any]]] = {}
+        self._objects: dict[UIElementId, UIElement[Any, Any]] = {}
         # mapping from object id to set of names that are bound to it
         self._bindings: dict[UIElementId, set[str]] = {}
         # mapping from object id to cell that created it
         self._constructing_cells: dict[UIElementId, CellId_t] = {}
+        # store of UI elements belonging to each cell
+        self._ui_elements_for_cell: dict[
+            CellId_t, list[UIElement[Any, Any]]
+        ] = {}
 
     def register(
         self,
@@ -43,9 +48,12 @@ def register(
             # its destructor was called, so manually delete the old element
             # here
             self.delete(object_id, id(self._objects[object_id]))
-        self._objects[object_id] = weakref.ref(ui_element)
+        self._objects[object_id] = ui_element
         assert execution_context is not None
         self._constructing_cells[object_id] = execution_context.cell_id
+        self._ui_elements_for_cell.setdefault(
+            execution_context.cell_id, []
+        ).append(ui_element)
         # bindings must be lazily registered, since there aren't any
         # bindings at UIElement object creation time
         if object_id in self._bindings:
@@ -66,8 +74,7 @@ def _has_parent_id(
         if child._id == parent_id:
             return True
         elif child._lens is not None:
-            element_ref = self._objects.get(child._lens.parent_id)
-            element = element_ref() if element_ref is not None else None
+            element = self._objects.get(child._lens.parent_id)
             if element is not None:
                 return self._has_parent_id(element, parent_id)
         return False
@@ -105,14 +112,7 @@ def _register_bindings(self, object_id: UIElementId) -> None:
     def get_object(self, object_id: UIElementId) -> UIElement[Any, Any]:
         if object_id not in self._objects:
             raise KeyError(f"UIElement with id {object_id} not found")
-        # UI elements are only updated if a global is bound to it. This ensures
-        # that the UI element update triggers reactivity, but also means that
-        # elements stored as, say, attributes on an object won't be updated.
-        if not self.bound_names(object_id):
-            raise NameError(f"UIElement with id {object_id} has no bindings")
-        obj = self._objects[object_id]()
-        assert obj is not None
-        return obj
+        return self._objects[object_id]
 
     def get_cell(self, object_id: UIElementId) -> CellId_t:
         return self._constructing_cells[object_id]
@@ -129,10 +129,7 @@ def resolve_lens(
         """
         if object_id not in self._objects:
             raise KeyError(f"UIElement with id {object_id} not found")
-        obj = self._objects[object_id]()
-        if obj is None:
-            raise RuntimeError(f"UIElement with id {object_id} was deleted")
-
+        obj = self._objects[object_id]
         lens = obj._lens
         if lens is None:
             # Base case: the element has no lens, so the resolved
@@ -153,7 +150,7 @@ def delete(self, object_id: UIElementId, python_id: int) -> None:
         if object_id not in self._objects:
             return
 
-        ui_element = self._objects[object_id]()
+        ui_element = self._objects[object_id]
         registered_python_id = (
             id(ui_element) if ui_element is not None else None
         )
@@ -172,3 +169,11 @@ def delete(self, object_id: UIElementId, python_id: int) -> None:
             del self._bindings[object_id]
         if object_id in self._constructing_cells:
             del self._constructing_cells[object_id]
+
+    def delete_elements_for_cell(self, cell_id: CellId_t) -> None:
+        if cell_id not in self._ui_elements_for_cell:
+            return
+
+        for element in self._ui_elements_for_cell[cell_id]:
+            self.delete(element._id, id(element))
+        del self._ui_elements_for_cell[cell_id]

marimo/_runtime/runtime.py

@@ -589,6 +589,7 @@ def _invalidate_cell_state(
         get_context().cell_lifecycle_registry.dispose(
             cell_id, deletion=deletion
         )
+        get_context().ui_element_registry.delete_elements_for_cell(cell_id)
         RemoveUIElements(cell_id=cell_id).broadcast()
 
     def _deactivate_cell(self, cell_id: CellId_t) -> set[CellId_t]:
@@ -1022,11 +1023,8 @@ async def set_ui_element_value(
                     object_id,
                     value,
                 )
-            except (KeyError, NameError):
-                # KeyError: A UI element may go out of scope if it was not
-                # assigned to a global variable
-                # NameError: UI element might not have bindings
-                LOGGER.debug("Could not find UIElement with id %s", object_id)
+            except KeyError:
+                LOGGER.error("Could not find UIElement with id %s", object_id)
                 continue
 
             with self._install_execution_context(