From 72d16f570e3bc06529b2ac8cbae72bea193510b6 Mon Sep 17 00:00:00 2001 From: Carson Sievert Date: Tue, 19 Nov 2024 15:57:13 -0600 Subject: [PATCH] Cleanup orphaned widget models (#167) --- CHANGELOG.md | 2 +- js/src/output.ts | 50 +++++++---- shinywidgets/_comm.py | 20 +++-- shinywidgets/_shinywidgets.py | 153 ++++++++++++++++++++++------------ shinywidgets/_utils.py | 13 +-- shinywidgets/static/output.js | 2 +- 6 files changed, 149 insertions(+), 91 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 72e56d2..aeace46 100644 --- 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 index 1ac98fc..6d97d68 100644 --- 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 index d7c7e22..fc0582a 100644 --- 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 index 649a905..1e711ce 100644 --- 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 index 7044434..25afd77 100644 --- 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 index 53ae1f2..2527dfd 100644 --- a/shinywidgets/static/output.js +++ b/shinywidgets/static/output.js @@ -36,7 +36,7 @@ eval("__webpack_require__.r(__webpack_exports__);\n/* harmony export */ __webpac \***********************/ /***/ ((__unused_webpack_module, __webpack_exports__, __webpack_require__) => { -eval("__webpack_require__.r(__webpack_exports__);\n/* harmony import */ var _jupyter_widgets_html_manager__WEBPACK_IMPORTED_MODULE_0__ = __webpack_require__(/*! @jupyter-widgets/html-manager */ \"@jupyter-widgets/html-manager\");\n/* harmony import */ var _jupyter_widgets_html_manager__WEBPACK_IMPORTED_MODULE_0___default = /*#__PURE__*/__webpack_require__.n(_jupyter_widgets_html_manager__WEBPACK_IMPORTED_MODULE_0__);\n/* harmony import */ var _comm__WEBPACK_IMPORTED_MODULE_1__ = __webpack_require__(/*! ./comm */ \"./src/comm.ts\");\n/* harmony import */ var _utils__WEBPACK_IMPORTED_MODULE_2__ = __webpack_require__(/*! ./utils */ \"./src/utils.ts\");\nvar _a;\n\n\n\n/******************************************************************************\n * Define a custom HTMLManager for use with Shiny\n ******************************************************************************/\nclass OutputManager extends _jupyter_widgets_html_manager__WEBPACK_IMPORTED_MODULE_0__.HTMLManager {\n // In a soon-to-be-released version of @jupyter-widgets/html-manager,\n // display_view()'s first \"dummy\" argument will be removed... this shim simply\n // makes it so that our manager can work with either version\n // https://github.com/jupyter-widgets/ipywidgets/commit/159bbe4#diff-45c126b24c3c43d2cee5313364805c025e911c4721d45ff8a68356a215bfb6c8R42-R43\n async display_view(view, options) {\n const n_args = super.display_view.length;\n if (n_args === 3) {\n return super.display_view({}, view, options);\n }\n else {\n // @ts-ignore\n return super.display_view(view, options);\n }\n }\n}\n// Define our own custom module loader for Shiny\nconst shinyRequireLoader = async function (moduleName, moduleVersion) {\n // shiny provides require.js and also sets `define.amd=false` to prevent