Use proper root_scope() API

Author: cpsievert

shinywidgets/_render_widget_base.py

@@ -70,7 +70,7 @@ def __init__(
         self._contexts: set[Context] = set()
 
     async def render(self) -> Jsonifiable | None:
-        with CurrentSessionOutput(self.output_id):
+        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)
@@ -216,20 +216,19 @@ def set_layout_defaults(widget: Widget) -> Tuple[Widget, bool]:
 
     return (widget, fill)
 
+class WidgetRenderContext:
+    """
+    Let the session when a widget is currently being rendered.
 
-# --------------------------------------------------------------------------------------------
-# 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:
+    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 = None
+        self._old_id = self.session.__dict__.get("__shinywidget_current_output_id")
 
     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
 
@@ -238,6 +237,6 @@ def __exit__(self, exc_type, exc_value, traceback):
         return False
 
     @staticmethod
-    def has_current_output(session):
+    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 CurrentSessionOutput, has_current_context
+from ._render_widget_base import WidgetRenderContext, has_current_context
 from ._utils import package_dir
 
 __all__ = (
@@ -60,9 +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
+    session = session.root_scope()
 
     # If this is the first time we've seen this session, initialize some things
     if session not in SESSIONS:
@@ -167,7 +165,7 @@ def on_close():
                 if id in WIDGET_INSTANCE_MAP:
                     del WIDGET_INSTANCE_MAP[id]
 
-        if CurrentSessionOutput.has_current_output(session):
+        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