Register cleanup methods when widget gets initialized in a reactive context

Author: cpsievert

js/src/comm.ts

@@ -96,8 +96,7 @@ export class ShinyComm {
     metadata?: any,
     buffers?: ArrayBuffer[] | ArrayBufferView[]
   ): string {
-    // Inform the server that a client-side model/comm has closed
-    Shiny.setInputValue("shinywidgets_comm_close", this.comm_id, {priority: "event"});
+    // I don't think we need to do anything here?
     return this.comm_id;
   }
 

js/src/output.ts

@@ -215,12 +215,6 @@ Shiny.addCustomMessageHandler("shinywidgets_comm_close", (msg_txt) => {
   });
 });
 
-// At least currently, all widgets must be created within a session scope, so we can
-// clear the state (i.e., .close() all the widget models) when the session ends
-$(document).on("shiny:disconnected", () => {
-  manager.clear_state();
-});
-
 // Our version of https://github.com/jupyter-widgets/widget-cookiecutter/blob/9694718/%7B%7Bcookiecutter.github_project_name%7D%7D/js/lib/extension.js#L8
 function setBaseURL(x: string = '') {
   const base_url = document.querySelector('body').getAttribute('data-base-url');

shinywidgets/_render_widget_base.py

@@ -70,9 +70,6 @@ def __init__(
         self._contexts: set[Context] = set()
 
     async def render(self) -> Jsonifiable | None:
-        # Before executing user function, get currently active widgets
-        widgets_before = set(WIDGET_INSTANCE_MAP.keys())
-
         value = await self.fn()
 
         # Attach value/widget attributes to user func so they can be accessed (in other reactive contexts)
@@ -91,16 +88,6 @@ async def render(self) -> Jsonifiable | None:
 
         self._widget = cast(WidgetT, widget)
 
-        # Collect widget ids introduced by the render function
-        widgets_after = set(WIDGET_INSTANCE_MAP.keys())
-        widgets_diff = widgets_after - widgets_before
-
-        # The next time this render context gets invalidated, close widgets introduced
-        # by the user function. This _could_ be problematic if widgets get hoisted
-        # out of the render function for use elsewhere, but if you're doing that,
-        # you're likely to have other problems anyway.
-        get_current_context().on_invalidate(lambda: close_widgets(widgets_diff))
-
         # Don't actually display anything unless this is a DOMWidget
         if not isinstance(widget, DOMWidget):
             return None
@@ -226,17 +213,3 @@ def set_layout_defaults(widget: Widget) -> Tuple[Widget, bool]:
             widget.chart = chart
 
     return (widget, fill)
-
-
-# Dictionary of all "active" widgets (ipywidgets automatically adds to this dictionary as
-# new widgets are created, but they won't get removed until the widget is explictly closed)
-WIDGET_INSTANCE_MAP = cast(dict[str, Widget], Widget.widgets)  # pyright: ignore[reportUnknownMemberType]
-
-# Given a set of widget ids, close those widgets. Closing should not only
-# remove the widget reference, but also send a message to the client to remove
-# the corresponding model from the widget manager.
-def close_widgets(ids: set[str]):
-    for id in ids:
-        widget = WIDGET_INSTANCE_MAP.get(id)
-        if widget is not None:
-            widget.close()

shinywidgets/_shinywidgets.py

@@ -26,7 +26,7 @@
 from ._cdn import SHINYWIDGETS_CDN_ONLY, SHINYWIDGETS_EXTENSION_WARNING
 from ._comm import BufferType, ShinyComm, ShinyCommManager
 from ._dependencies import require_dependency
-from ._render_widget_base import WIDGET_INSTANCE_MAP
+from ._render_widget_base import has_current_context
 from ._utils import is_instance_of_class, package_dir
 
 __all__ = (
@@ -79,12 +79,16 @@ def init_shiny_widget(w: Widget):
 
     # By the time we get here, the user has already had an opportunity to specify a model_id,
     # so it isn't yet populated, generate a random one so we can assign the same id to the comm
-    if getattr(w, "_model_id", None) is None:
-        w._model_id = uuid4().hex
+    w._model_id = getattr(w, "_model_id", uuid4().hex)
+
+    if not isinstance(w._model_id, str):
+        raise ValueError(
+            f"Expected widget id to be a string, but got {type(w._model_id)}"
+        )
 
     # Initialize the comm...this will also send the initial state of the widget
     w.comm = ShinyComm(
-        comm_id=w._model_id,  # pyright: ignore
+        comm_id=w._model_id,
         comm_manager=COMM_MANAGER,
         target_name="jupyter.widgets",
         data={"state": state, "buffer_paths": buffer_paths},
@@ -94,6 +98,14 @@ def init_shiny_widget(w: Widget):
         html_deps=session._process_ui(TagList(widget_dep))["deps"],
     )
 
+    # If we're in a reactive context, close this widget when the context is invalidated
+    if has_current_context():
+        ctx = get_current_context()
+        ctx.on_invalidate(lambda: w.close())
+
+    # Remember to close the widget when the session ends
+    session.on_ended(lambda: w.close())
+
     # Some widget's JS make external requests for static files (e.g.,
     # ipyleaflet markers) under this resource path. Note that this assumes that
     # we're setting the data-base-url attribute on the <body> (which we should
@@ -111,7 +123,7 @@ def init_shiny_widget(w: Widget):
     # everything after this point should be done once per session
     if session in SESSIONS:
         return
-    SESSIONS.add(session)  # type: ignore
+    SESSIONS.add(session)
 
     # Somewhere inside ipywidgets, it makes requests for static files
     # under the publicPath set by the webpack.config.js file.
@@ -133,21 +145,10 @@ def _():
             comm: ShinyComm = COMM_MANAGER.comms[comm_id]
             comm.handle_msg(msg)
 
-    # Handle a close message from the client.
-    @reactive.effect
-    @reactive.event(session.input.shinywidgets_comm_close)
-    def _():
-        comm_id = session.input.shinywidgets_comm_close()
-        # Close the widget, which unregisters/deletes the comm, and also drops
-        # ipywidget's reference to the instance, allowing it to be garbage collected.
-        w_obj = WIDGET_INSTANCE_MAP.get(comm_id)
-        if w_obj:
-            w_obj.close()
-
     def _restore_state():
         if old_comm_klass is not None:
             Widget.comm.klass = old_comm_klass  # type: ignore
-        SESSIONS.remove(session)  # type: ignore
+        SESSIONS.remove(session)
 
     session.on_ended(_restore_state)
 
@@ -156,7 +157,7 @@ def _restore_state():
 Widget.on_widget_constructed(init_shiny_widget)  # type: ignore
 
 # Use WeakSet() over Set() so that the session can be garbage collected
-SESSIONS = WeakSet()  # type: ignore
+SESSIONS: WeakSet[Session] = WeakSet()
 COMM_MANAGER = ShinyCommManager()