Only cleanup widgets that get initialized in an output context (#191)

Author: cpsievert

CHANGELOG.md

@@ -5,6 +5,10 @@ All notable changes to shinywidgets will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [UNRELEASED]
+
+* Widgets initialized inside a `reactive.effect()` are no longer automatically removed when the effect invalidates. (#191)
+
 ## [0.5.2]  - 2025-04-04
 
 * Constructing a widget inside of a `shiny.reactive.ExtendedTask()` no longer errors out. (#188)

shinywidgets/_render_widget_base.py

@@ -11,6 +11,7 @@
     get_current_context,  # pyright: ignore[reportPrivateImportUsage]
 )
 from shiny.render.renderer import Jsonifiable, Renderer, ValueFn
+from shiny.session import require_active_session
 
 from ._as_widget import as_widget
 from ._dependencies import widget_pkg
@@ -69,7 +70,8 @@ def __init__(
         self._contexts: set[Context] = set()
 
     async def render(self) -> Jsonifiable | None:
-        value = await self.fn()
+        with WidgetRenderContext(self.output_id):
+            value = await self.fn()
 
         # Attach value/widget attributes to user func so they can be accessed (in other reactive contexts)
         self._value = value
@@ -213,3 +215,28 @@ def set_layout_defaults(widget: Widget) -> Tuple[Widget, bool]:
             widget.chart = chart
 
     return (widget, fill)
+
+class WidgetRenderContext:
+    """
+    Let the session when a widget is currently being rendered.
+
+    This is used to ensure that widget's that are initialized in a render_widget()
+    context are cleaned up properly when that context is re-entered.
+    """
+    def __init__(self, output_id):
+        self.session = require_active_session(None)
+        self.output_id = output_id
+        self._old_id = self.session.__dict__.get("__shinywidget_current_output_id")
+
+    def __enter__(self):
+        self.session.__dict__["__shinywidget_current_output_id"] = self.output_id
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.session.__dict__["__shinywidget_current_output_id"] = self._old_id
+        return False
+
+    @staticmethod
+    def is_rendering_widget(session):
+        id = session.__dict__.get("__shinywidget_current_output_id")
+        return id is not None

shinywidgets/_shinywidgets.py

@@ -29,7 +29,7 @@
 from ._cdn import SHINYWIDGETS_CDN_ONLY, SHINYWIDGETS_EXTENSION_WARNING
 from ._comm import BufferType, OrphanedShinyComm, ShinyComm, ShinyCommManager
 from ._dependencies import require_dependency
-from ._render_widget_base import has_current_context
+from ._render_widget_base import WidgetRenderContext, has_current_context
 from ._utils import package_dir
 
 __all__ = (
@@ -60,8 +60,7 @@ def init_shiny_widget(w: Widget):
         return
     # Break out of any module-specific session. Otherwise, input.shinywidgets_comm_send
     # will be some module-specific copy.
-    while hasattr(session, "_parent"):
-        session = cast(Session, session._parent)  # pyright: ignore
+    session = session.root_scope()
 
     # If this is the first time we've seen this session, initialize some things
     if session not in SESSIONS:
@@ -148,12 +147,8 @@ def _open_shiny_comm():
 
         _open_shiny_comm.destroy()
 
-    # If we're in a reactive context, close this widget when the context is invalidated
-    # TODO: this should probably only be done in an output context, but I'm pretty sure
-    # we don't have a decent way to determine that at the moment. In theory, doing this
-    # in _any_ reactive context be problematic if you have an effect() that adds one
-    # widget to another (i.e., a marker to a map) and want that marker to persist through
-    # the next invalidation. The example provided in #174 is one such example.
+    # If the widget initialized in a reactive _output_ context, then cleanup the widget
+    # when the context gets invalidated.
     if has_current_context():
         ctx = get_current_context()
 
@@ -170,13 +165,7 @@ def on_close():
                 if id in WIDGET_INSTANCE_MAP:
                     del WIDGET_INSTANCE_MAP[id]
 
-        # This could be running in a shiny.reactive.ExtendedTask, in which case,
-        # the context is a DenialContext. As a result, on_invalidate() will throw
-        # (since reading/invalidating reactive sources isn't allowed in this context).
-        # For now, we just don't clean up the widget in this case.
-        # TODO: this line can likely be removed once we start closing iff we're in a
-        # output context (see TODO comment above)
-        if "DenialContext" != ctx.__class__.__name__:
+        if WidgetRenderContext.is_rendering_widget(session):
             ctx.on_invalidate(on_close)
 
     # Keep track of what session this widget belongs to (so we can close it when the