Cleanup orphaned widget models (#167)

cpsievert · web-flow · commit 72d16f570e3b · 2024-11-19T15:57:13.000-06:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -7,7 +7,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [UNRELEASED]
 
-
+* Fixed a memory leak issue. (#167)
 
 ## [0.3.4] - 2024-10-29
 
diff --git a/js/src/output.ts b/js/src/output.ts
@@ -100,13 +100,6 @@ class IPyWidgetOutput extends Shiny.OutputBinding {
     const view = await manager.create_view(model, {});
     await manager.display_view(view, {el: el});
 
-    // Don't allow more than one .lmWidget container, which can happen
-    // when the view is displayed more than once
-    // TODO: It's probably better to get view(s) from m.views and .remove() them
-    while (el.childNodes.length > 1) {
-      el.removeChild(el.childNodes[0]);
-    }
-
     // The ipywidgets container (.lmWidget)
     const lmWidget = el.children[0] as HTMLElement;
 
@@ -189,21 +182,46 @@ Shiny.addCustomMessageHandler("shinywidgets_comm_open", (msg_txt) => {
 // Basically out version of https://github.com/jupyterlab/jupyterlab/blob/d33de15/packages/services/src/kernel/default.ts#L1200-L1215
 Shiny.addCustomMessageHandler("shinywidgets_comm_msg", (msg_txt) => {
   const msg = jsonParse(msg_txt);
-  manager.get_model(msg.content.comm_id).then(m => {
-    // @ts-ignore for some reason IClassicComm doesn't have this method, but we do
-    m.comm.handle_msg(msg);
-  });
+  const id = msg.content.comm_id;
+  const model = manager.get_model(id);
+  if (!model) {
+    console.error(`Couldn't handle message for model ${id} because it doesn't exist.`);
+    return;
+  }
+  model
+    .then(m => {
+      // @ts-ignore for some reason IClassicComm doesn't have this method, but we do
+      m.comm.handle_msg(msg);
+    })
+    .catch(console.error);
 });
 
-// TODO: test that this actually works
+
+// Handle the closing of a widget/comm/model
 Shiny.addCustomMessageHandler("shinywidgets_comm_close", (msg_txt) => {
   const msg = jsonParse(msg_txt);
-  manager.get_model(msg.content.comm_id).then(m => {
-    // @ts-ignore for some reason IClassicComm doesn't have this method, but we do
-    m.comm.handle_close(msg)
-  });
+  const id = msg.content.comm_id;
+  const model = manager.get_model(id);
+  if (!model) {
+    console.error(`Couldn't close model ${id} because it doesn't exist.`);
+    return;
+  }
+  model
+    .then(m => {
+        // Closing the model removes the corresponding view from the DOM
+        m.close();
+        // .close() isn't enough to remove manager's reference to it,
+        // and apparently the only way to remove it is through the `comm:close` event
+        // https://github.com/jupyter-widgets/ipywidgets/blob/303cae4/packages/base-manager/src/manager-base.ts#L330-L337
+        // https://github.com/jupyter-widgets/ipywidgets/blob/303cae4/packages/base/src/widget.ts#L251-L253
+        m.trigger("comm:close");
+    })
+    .catch(console.error);
 });
 
+$(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 = '') {
diff --git a/shinywidgets/_comm.py b/shinywidgets/_comm.py
@@ -97,9 +97,10 @@ def close(
             return
         self._closed = True
         data = self._closed_data if data is None else data
-        self._publish_msg(
-            "shinywidgets_comm_close", data=data, metadata=metadata, buffers=buffers
-        )
+        if get_current_session():
+            self._publish_msg(
+                "shinywidgets_comm_close", data=data, metadata=metadata, buffers=buffers
+            )
         if not deleting:
             # If deleting, the comm can't be unregistered
             self.comm_manager.unregister_comm(self)
@@ -169,10 +170,17 @@ def _publish_msg(
         def _send():
             run_coro_hybrid(session.send_custom_message(msg_type, msg_txt))  # type: ignore
 
-        # N.B., if we don't do this on flush, then if you initialize a widget
-        # outside of a reactive context, run_coro_sync() will complain with
+        # N.B., if messages are sent immediately, run_coro_sync() could fail with
         # 'async function yielded control; it did not finish in one iteration.'
-        session.on_flush(_send)
+        # if executed outside of a reactive context.
+        if msg_type == "shinywidgets_comm_close":
+            # The primary way widgets are closed are when a new widget is rendered in
+            # its place (see render_widget_base). By sending close on_flushed(), we
+            # ensure to close the 'old' widget after the new one is created. (avoiding a
+            # "flicker" of the old widget being removed before the new one is created)
+            session.on_flushed(_send)
+        else:
+            session.on_flush(_send)
 
     # This is the method that ipywidgets.widgets.Widget uses to respond to client-side changes
     def on_msg(self, callback: MsgCallback) -> None:
diff --git a/shinywidgets/_shinywidgets.py b/shinywidgets/_shinywidgets.py
@@ -3,7 +3,8 @@
 import copy
 import json
 import os
-from typing import Any, Optional, Sequence, Union, cast
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Any, Optional, Sequence, TypeGuard, Union, cast
 from uuid import uuid4
 from weakref import WeakSet
 
@@ -26,13 +27,17 @@
 from ._cdn import SHINYWIDGETS_CDN_ONLY, SHINYWIDGETS_EXTENSION_WARNING
 from ._comm import BufferType, ShinyComm, ShinyCommManager
 from ._dependencies import require_dependency
-from ._utils import is_instance_of_class, package_dir
+from ._render_widget_base import has_current_context
+from ._utils import package_dir
 
 __all__ = (
     "register_widget",
     "reactive_read",
 )
 
+if TYPE_CHECKING:
+    from traitlets.traitlets import Instance
+
 
 # --------------------------------------------------------------------------------------------
 # When a widget is initialized, also initialize a communication channel (via the Shiny
@@ -54,19 +59,43 @@ def init_shiny_widget(w: Widget):
     while hasattr(session, "_parent"):
         session = cast(Session, session._parent)  # pyright: ignore
 
-    # Previous versions of ipywidgets (< 8.0.5) had
-    #   `Widget.comm = Instance('ipykernel.comm.Comm')`
-    # which meant we'd get a runtime error when setting `Widget.comm = ShinyComm()`.
-    # In more recent versions, this is no longer necessary since they've (correctly)
-    # changed comm from an Instance() to Any().
-    # https://github.com/jupyter-widgets/ipywidgets/pull/3533/files#diff-522bb5e7695975cba0199c6a3d6df5be827035f4dc18ed6da22ac216b5615c77R482
-    old_comm_klass = None
-    if is_instance_of_class(Widget.comm, "Instance", "traitlets.traitlets"):  # type: ignore
-        old_comm_klass = copy.copy(Widget.comm.klass)  # type: ignore
-        Widget.comm.klass = object  # type: ignore
+    # If this is the first time we've seen this session, initialize some things
+    if session not in SESSIONS:
+        SESSIONS.add(session)
+
+        # Somewhere inside ipywidgets, it makes requests for static files
+        # under the publicPath set by the webpack.config.js file.
+        session.app._dependency_handler.mount(
+            "/dist/",
+            StaticFiles(directory=os.path.join(package_dir("shinywidgets"), "static")),
+            name="shinywidgets-static-resources",
+        )
+
+        # Handle messages from the client. Note that widgets like qgrid send client->server messages
+        # to figure out things like what filter to be shown in the table.
+        @reactive.effect
+        @reactive.event(session.input.shinywidgets_comm_send)
+        def _():
+            msg_txt = session.input.shinywidgets_comm_send()
+            msg = json.loads(msg_txt)
+            comm_id = msg["content"]["comm_id"]
+            if comm_id in COMM_MANAGER.comms:
+                comm: ShinyComm = COMM_MANAGER.comms[comm_id]
+                comm.handle_msg(msg)
+
+        def _cleanup_session_state():
+            SESSIONS.remove(session)
+            # Cleanup any widgets that were created in this session
+            for id in SESSION_WIDGET_ID_MAP[session.id]:
+                widget = WIDGET_INSTANCE_MAP.get(id)
+                if widget:
+                    widget.close()
+            del SESSION_WIDGET_ID_MAP[session.id]
+
+        session.on_ended(_cleanup_session_state)
 
     # Get the initial state of the widget
-    state, buffer_paths, buffers = _remove_buffers(w.get_state())  # type: ignore
+    state, buffer_paths, buffers = _remove_buffers(w.get_state())
 
     # Make sure window.require() calls made by 3rd party widgets
     # (via handle_comm_open() -> new_model() -> loadClass() -> requireLoader())
@@ -81,17 +110,29 @@ def init_shiny_widget(w: Widget):
     if getattr(w, "_model_id", None) is None:
         w._model_id = uuid4().hex
 
+    id = cast(str, w._model_id)  # pyright: ignore[reportUnknownMemberType]
+
     # Initialize the comm...this will also send the initial state of the widget
-    w.comm = ShinyComm(
-        comm_id=w._model_id,  # pyright: ignore
-        comm_manager=COMM_MANAGER,
-        target_name="jupyter.widgets",
-        data={"state": state, "buffer_paths": buffer_paths},
-        buffers=cast(BufferType, buffers),
-        # TODO: should this be hard-coded?
-        metadata={"version": __protocol_version__},
-        html_deps=session._process_ui(TagList(widget_dep))["deps"],
-    )
+    with widget_comm_patch():
+        w.comm = ShinyComm(
+            comm_id=id,
+            comm_manager=COMM_MANAGER,
+            target_name="jupyter.widgets",
+            data={"state": state, "buffer_paths": buffer_paths},
+            buffers=cast(BufferType, buffers),
+            # TODO: should this be hard-coded?
+            metadata={"version": __protocol_version__},
+            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())
+
+    # Keep track of what session this widget belongs to (so we can close it when the
+    # session ends)
+    SESSION_WIDGET_ID_MAP.setdefault(session.id, []).append(id)
 
     # Some widget's JS make external requests for static files (e.g.,
     # ipyleaflet markers) under this resource path. Note that this assumes that
@@ -107,45 +148,21 @@ def init_shiny_widget(w: Widget):
                 name=f"{widget_dep.name}-nbextension-static-resources",
             )
 
-    # everything after this point should be done once per session
-    if session in SESSIONS:
-        return
-    SESSIONS.add(session)  # type: ignore
-
-    # Somewhere inside ipywidgets, it makes requests for static files
-    # under the publicPath set by the webpack.config.js file.
-    session.app._dependency_handler.mount(
-        "/dist/",
-        StaticFiles(directory=os.path.join(package_dir("shinywidgets"), "static")),
-        name="shinywidgets-static-resources",
-    )
-
-    # Handle messages from the client. Note that widgets like qgrid send client->server messages
-    # to figure out things like what filter to be shown in the table.
-    @reactive.Effect
-    @reactive.event(session.input.shinywidgets_comm_send)
-    def _():
-        msg_txt = session.input.shinywidgets_comm_send()
-        msg = json.loads(msg_txt)
-        comm_id = msg["content"]["comm_id"]
-        comm: ShinyComm = COMM_MANAGER.comms[comm_id]
-        comm.handle_msg(msg)
-
-    def _restore_state():
-        if old_comm_klass is not None:
-            Widget.comm.klass = old_comm_klass  # type: ignore
-        SESSIONS.remove(session)  # type: ignore
-
-    session.on_ended(_restore_state)
-
 
 # TODO: can we restore the widget constructor in a sensible way?
 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()
 
+# Dictionary mapping session id to widget ids
+# The key is the session id, and the value is a list of widget ids
+SESSION_WIDGET_ID_MAP: dict[str, list[str]] = {}
+
+# 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]
 
 # --------------------------------------
 # Reactivity
@@ -216,6 +233,32 @@ def _():
 
     return w
 
+# Previous versions of ipywidgets (< 8.0.5) had
+#   `Widget.comm = Instance('ipykernel.comm.Comm')`
+# which meant we'd get a runtime error when setting `Widget.comm = ShinyComm()`.
+# In more recent versions, this is no longer necessary since they've (correctly)
+# changed comm from an Instance() to Any().
+# https://github.com/jupyter-widgets/ipywidgets/pull/3533/files#diff-522bb5e7695975cba0199c6a3d6df5be827035f4dc18ed6da22ac216b5615c77R482
+@contextmanager
+def widget_comm_patch():
+    if not is_traitlet_instance(Widget.comm):
+        yield
+        return
+
+    comm_klass = copy.copy(Widget.comm.klass)
+    Widget.comm.klass = object
+
+    yield
+
+    Widget.comm.klass = comm_klass
+
+
+def is_traitlet_instance(x: object) -> "TypeGuard[Instance]":
+    try:
+        from traitlets.traitlets import Instance
+    except ImportError:
+        return False
+    return isinstance(x, Instance)
 
 # It doesn't, at the moment, seem feasible to establish a comm with statically rendered widgets,
 # and partially for this reason, it may not be sensible to provide an input-like API for them.
diff --git a/shinywidgets/_utils.py b/shinywidgets/_utils.py
@@ -1,7 +1,7 @@
 import importlib
 import os
 import tempfile
-from typing import Optional
+
 
 # similar to base::system.file()
 def package_dir(package: str) -> str:
@@ -10,14 +10,3 @@ def package_dir(package: str) -> str:
         if pkg_file is None:
             raise ImportError(f"Couldn't load package {package}")
         return os.path.dirname(pkg_file)
-
-
-def is_instance_of_class(
-    x: object, class_name: str, module_name: Optional[str] = None
-) -> bool:
-    typ = type(x)
-    res = typ.__name__ == class_name
-    if module_name is None:
-        return res
-    else:
-        return res and typ.__module__ == module_name
diff --git a/shinywidgets/static/output.js b/shinywidgets/static/output.js
