Skip to content

Commit

Permalink
Cleanup orphaned widget models (#167)
Browse files Browse the repository at this point in the history
  • Loading branch information
cpsievert authored Nov 19, 2024
1 parent d625c3f commit 72d16f5
Showing 6 changed files with 149 additions and 91 deletions.
2 changes: 1 addition & 1 deletion CHANGELOG.md
Original file line number Diff line number Diff line change
@@ -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

50 changes: 34 additions & 16 deletions js/src/output.ts
Original file line number Diff line number Diff line change
@@ -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 = '') {
20 changes: 14 additions & 6 deletions shinywidgets/_comm.py
Original file line number Diff line number Diff line change
@@ -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:
153 changes: 98 additions & 55 deletions shinywidgets/_shinywidgets.py
Original file line number Diff line number Diff line change
@@ -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.
13 changes: 1 addition & 12 deletions shinywidgets/_utils.py
Original file line number Diff line number Diff line change
@@ -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
Loading

0 comments on commit 72d16f5

Please sign in to comment.