From 244f25f23bc278a80be847d4fddacb9da29319c4 Mon Sep 17 00:00:00 2001 From: Akshay Agrawal Date: Sat, 11 May 2024 09:51:50 -0700 Subject: [PATCH] improvement: hold on to unnamed UI elements Changes the semantics of the UIRegistry to hold on to unnamed UI elements. This makes function call requests work on unnamed UI elements. The registry now keeps track of all the UI elements registered to a cell, and exposes a method to unregister them. The runtime unregisters all UI elements for a cell when cell state is invalidated. --- marimo/_plugins/ui/_core/registry.py | 45 +++++++++++++++------------- marimo/_runtime/runtime.py | 8 ++--- 2 files changed, 28 insertions(+), 25 deletions(-) diff --git a/marimo/_plugins/ui/_core/registry.py b/marimo/_plugins/ui/_core/registry.py index 955e506b01..a6427648b7 100644 --- a/marimo/_plugins/ui/_core/registry.py +++ b/marimo/_plugins/ui/_core/registry.py @@ -2,8 +2,7 @@ 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 @@ -11,10 +10,12 @@ 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] diff --git a/marimo/_runtime/runtime.py b/marimo/_runtime/runtime.py index 8d28ee5b1d..55c2b35f4e 100644 --- a/marimo/_runtime/runtime.py +++ b/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(