Close #190: Only cleanup widgets that get initialized in an output context

Author: cpsievert

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 CurrentSessionOutput(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,29 @@ def set_layout_defaults(widget: Widget) -> Tuple[Widget, bool]:
             widget.chart = chart
 
     return (widget, fill)
+
+
+# --------------------------------------------------------------------------------------------
+# Context manager to set the current output id
+# (this is needed since, in order to clean up widgets properly, we need to
+#  know if they were initialized in a output context or not)
+# --------------------------------------------------------------------------------------------
+class CurrentSessionOutput:
+    def __init__(self, output_id):
+        self.session = require_active_session(None)
+        self.output_id = output_id
+        self._old_id = None
+
+    def __enter__(self):
+        self._old_id = self.session.__dict__.get("__shinywidget_current_output_id")
+        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 has_current_output(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 CurrentSessionOutput, has_current_context
 from ._utils import package_dir
 
 __all__ = (
@@ -60,6 +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.
+    # TODO: session = session.root_scope()
     while hasattr(session, "_parent"):
         session = cast(Session, session._parent)  # pyright: ignore
 
@@ -148,12 +149,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 +167,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 CurrentSessionOutput.has_current_output(session):
             ctx.on_invalidate(on_close)
 
     # Keep track of what session this widget belongs to (so we can close it when the