feat(experimental): Add command system for calling Python from JS (#453)

Co-authored-by: Talley Lambert <talley.lambert@gmail.com>

Author: manzt

.changeset/nervous-goats-rescue.md

@@ -0,0 +1,28 @@
+---
+"anywidget": patch
+"@anywidget/types": patch
+---
+
+Add experimental `invoke` API to call Python functions from the front end and
+await the response.
+
+This removes a lot of boilerplate required for this pattern. The API is
+experimental and opt-in only. Subclasses must use the `command` to register
+functions.
+
+```py
+class Widget(anywidget.AnyWidget):
+    _esm = """
+    export default {
+      async render({ model, el, experimental }) {
+        let [msg, buffers] = await experimental.invoke("_echo", "hello, world");
+        console.log(msg); // "HELLO, WORLD"
+      },
+    };
+    """
+
+    @anywidget.experimental.command
+    def _echo(self, msg, buffers):
+        # upper case the message
+        return msg.upper(), buffers
+```

anywidget/_descriptor.py

@@ -57,8 +57,7 @@
     from ._protocols import CommMessage
 
     class _GetState(Protocol):
-        def __call__(self, obj: Any, include: set[str] | None) -> dict:
-            ...
+        def __call__(self, obj: Any, include: set[str] | None) -> dict: ...
 
     # catch all for types that can be serialized ... too hard to actually type
     Serializable: TypeAlias = Any
@@ -205,12 +204,10 @@ def __set_name__(self, owner: type, name: str) -> None:
         self._name = name
 
     @overload
-    def __get__(self, instance: None, owner: type) -> MimeBundleDescriptor:
-        ...
+    def __get__(self, instance: None, owner: type) -> MimeBundleDescriptor: ...
 
     @overload
-    def __get__(self, instance: object, owner: type) -> ReprMimeBundle:
-        ...
+    def __get__(self, instance: object, owner: type) -> ReprMimeBundle: ...
 
     def __get__(
         self, instance: object | None, owner: type

anywidget/_protocols.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Sequence
+from typing import TYPE_CHECKING, Any, Callable, Sequence
 
 from typing_extensions import Literal, Protocol, TypedDict
 
@@ -40,7 +40,7 @@ class CommMessage(TypedDict):
 
 
 class MimeReprCallable(Protocol):
-    """Protocol for _repr_mimebundle.
+    """Protocol for _repr_mimebundle_.
 
     https://ipython.readthedocs.io/en/stable/config/integrating.html#more-powerful-methods
 
@@ -52,11 +52,20 @@ class MimeReprCallable(Protocol):
 
     def __call__(
         self, include: Sequence[str], exclude: Sequence[str]
-    ) -> dict | tuple[dict, dict]:
-        ...
+    ) -> dict | tuple[dict, dict]: ...
 
 
 class AnywidgetProtocol(Protocol):
     """Anywidget classes have a MimeBundleDescriptor at `_repr_mimebundle_`."""
 
     _repr_mimebundle_: MimeBundleDescriptor
+
+
+class WidgetBase(Protocol):
+    """Widget subclasses with a custom message reducer."""
+
+    def send(self, msg: str | dict | list, buffers: list[bytes]) -> None: ...
+
+    def on_msg(
+        self, callback: Callable[[Any, str | list | dict, list[bytes]], None]
+    ) -> None: ...

anywidget/experimental.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import contextlib
 import dataclasses
 import pathlib
 import typing
@@ -8,6 +9,9 @@
 
 from ._descriptor import MimeBundleDescriptor
 
+if typing.TYPE_CHECKING:  # pragma: no cover
+    from ._protocols import WidgetBase
+
 __all__ = ["dataclass", "widget", "MimeBundleDescriptor"]
 
 _T = typing.TypeVar("_T")
@@ -103,3 +107,62 @@ def _decorator(cls: T) -> T:
         return cls
 
     return _decorator(cls) if cls is not None else _decorator  # type: ignore
+
+
+_ANYWIDGET_COMMAND = "_anywidget_command"
+
+_AnyWidgetCommand = typing.Callable[
+    [typing.Any, typing.Any, typing.List[bytes]],
+    typing.Tuple[typing.Any, typing.List[bytes]],
+]
+
+
+def command(cmd: _AnyWidgetCommand) -> _AnyWidgetCommand:
+    """Mark a function as a command for anywidget."""
+    setattr(cmd, _ANYWIDGET_COMMAND, True)
+    return cmd
+
+
+_AnyWidgetCommandBound = typing.Callable[
+    [typing.Any, typing.List[bytes]], typing.Tuple[typing.Any, typing.List[bytes]]
+]
+
+
+def _collect_commands(widget: WidgetBase) -> dict[str, _AnyWidgetCommandBound]:
+    cmds: dict[str, _AnyWidgetCommandBound] = {}
+    for attr_name in dir(widget):
+        # suppressing silly assertion erro from ipywidgets _staticproperty
+        # ref: https://github.com/jupyter-widgets/ipywidgets/blob/b78de43e12ff26e4aa16e6e4c6844a7c82a8ee1c/python/ipywidgets/ipywidgets/widgets/widget.py#L291-L297
+        with contextlib.suppress(AssertionError):
+            attr = getattr(widget, attr_name)
+            if callable(attr) and getattr(attr, _ANYWIDGET_COMMAND, False):
+                cmds[attr_name] = attr
+    return cmds
+
+
+def _register_anywidget_commands(
+    widget: WidgetBase,
+) -> None:
+    """Register a custom message reducer for a widget if it implements the protocol."""
+    # Only add the callback if the widget has any commands.
+    cmds = _collect_commands(widget)
+    if len(cmds) == 0:
+        return None
+
+    def handle_anywidget_command(
+        self: WidgetBase, msg: str | list | dict, buffers: list[bytes]
+    ) -> None:
+        if not isinstance(msg, dict) or msg.get("kind") != "anywidget-command":
+            return
+        cmd = cmds[msg["name"]]
+        response, buffers = cmd(msg["msg"], buffers)
+        self.send(
+            {
+                "id": msg["id"],
+                "kind": "anywidget-command-response",
+                "response": response,
+            },
+            buffers,
+        )
+
+    widget.on_msg(handle_anywidget_command)

anywidget/widget.py

@@ -17,6 +17,7 @@
     try_file_contents,
 )
 from ._version import __version__
+from .experimental import _register_anywidget_commands
 
 
 class AnyWidget(ipywidgets.DOMWidget):  # type: ignore [misc]
@@ -57,6 +58,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
 
         self.add_traits(**anywidget_traits)
         super().__init__(*args, **kwargs)
+        _register_anywidget_commands(self)
 
     def __init_subclass__(cls, **kwargs: dict) -> None:
         """Coerces _esm and _css to FileContents if they are files."""

packages/anywidget/package.json

@@ -24,6 +24,7 @@
 		"@anywidget/types": "workspace:~",
 		"@anywidget/vite": "workspace:~",
 		"@jupyter-widgets/base": "^2 || ^3 || ^4 || ^5 || ^6",
+		"@lukeed/uuid": "^2.0.1",
 		"solid-js": "^1.8.16"
 	},
 	"devDependencies": {

packages/anywidget/src/widget.js

@@ -1,9 +1,5 @@
-import {
-	createEffect,
-	createResource,
-	createRoot,
-	createSignal,
-} from "solid-js";
+import * as uuid from "@lukeed/uuid";
+import * as solid from "solid-js";
 
 /**
  * @typedef AnyWidget
@@ -242,6 +238,50 @@ function throw_anywidget_error(source) {
 	throw source;
 }
 
+/**
+ * @template T
+ * @param {import("@anywidget/types").AnyModel} model
+ * @param {string} name
+ * @param {any} [msg]
+ * @param {DataView[]} [buffers]
+ * @param {{ timeout?: number }} [options]
+ * @return {Promise<[T, DataView[]]>}
+ */
+export function invoke(
+	model,
+	name,
+	msg,
+	buffers = [],
+	{ timeout = 3000 } = {},
+) {
+	// crypto.randomUUID() is not available in non-secure contexts (i.e., http://)
+	// so we use simple (non-secure) polyfill.
+	let id = uuid.v4();
+	return new Promise((resolve, reject) => {
+		let timer = setTimeout(() => {
+			reject(new Error(`Promise timed out after ${timeout} ms`));
+			model.off("msg:custom", handler);
+		}, timeout);
+
+		/**
+		 * @param {{ id: string, kind: "anywidget-command-response", response: T }} msg
+		 * @param {DataView[]} buffers
+		 */
+		function handler(msg, buffers) {
+			if (!(msg.id === id)) return;
+			clearTimeout(timer);
+			resolve([msg.response, buffers]);
+			model.off("msg:custom", handler);
+		}
+		model.on("msg:custom", handler);
+		model.send(
+			{ id, kind: "anywidget-command", name, msg },
+			undefined,
+			buffers,
+		);
+	});
+}
+
 class Runtime {
 	/** @type {() => void} */
 	#disposer = () => {};
@@ -253,34 +293,38 @@ class Runtime {
 
 	/** @param {import("@jupyter-widgets/base").DOMWidgetModel} model */
 	constructor(model) {
-		this.#disposer = createRoot((dispose) => {
-			let [css, set_css] = createSignal(model.get("_css"));
+		this.#disposer = solid.createRoot((dispose) => {
+			let [css, set_css] = solid.createSignal(model.get("_css"));
 			model.on("change:_css", () => {
 				let id = model.get("_anywidget_id");
 				console.debug(`[anywidget] css hot updated: ${id}`);
 				set_css(model.get("_css"));
 			});
-			createEffect(() => {
+			solid.createEffect(() => {
 				let id = model.get("_anywidget_id");
 				load_css(css(), id);
 			});
 
 			/** @type {import("solid-js").Signal<string>} */
-			let [esm, setEsm] = createSignal(model.get("_esm"));
+			let [esm, setEsm] = solid.createSignal(model.get("_esm"));
 			model.on("change:_esm", async () => {
 				let id = model.get("_anywidget_id");
 				console.debug(`[anywidget] esm hot updated: ${id}`);
 				setEsm(model.get("_esm"));
 			});
 			/** @type {void | (() => import("vitest").Awaitable<void>)} */
 			let cleanup;
-			this.#widget_result = createResource(esm, async (update) => {
+			this.#widget_result = solid.createResource(esm, async (update) => {
 				await safe_cleanup(cleanup, "initialize");
 				try {
 					model.off(null, null, INITIALIZE_MARKER);
 					let widget = await load_widget(update);
 					cleanup = await widget.initialize?.({
 						model: model_proxy(model, INITIALIZE_MARKER),
+						experimental: {
+							// @ts-expect-error - bind isn't working
+							invoke: invoke.bind(null, model),
+						},
 					});
 					return ok(widget);
 				} catch (e) {
@@ -302,11 +346,11 @@ class Runtime {
 	 */
 	async create_view(view) {
 		let model = view.model;
-		let disposer = createRoot((dispose) => {
+		let disposer = solid.createRoot((dispose) => {
 			/** @type {void | (() => import("vitest").Awaitable<void>)} */
 			let cleanup;
 			let resource =
-				createResource(this.#widget_result, async (widget_result) => {
+				solid.createResource(this.#widget_result, async (widget_result) => {
 					cleanup?.();
 					// Clear all previous event listeners from this hook.
 					model.off(null, null, view);
@@ -319,12 +363,16 @@ class Runtime {
 						cleanup = await widget.render?.({
 							model: model_proxy(model, view),
 							el: view.el,
+							experimental: {
+								// @ts-expect-error - bind isn't working
+								invoke: invoke.bind(null, model),
+							},
 						});
 					} catch (e) {
 						throw_anywidget_error(e);
 					}
 				})[0];
-			createEffect(() => {
+			solid.createEffect(() => {
 				if (resource.error) {
 					// TODO: Show error in the view?
 				}

packages/types/index.ts

@@ -43,9 +43,18 @@ export interface AnyModel<T extends ObjectHash = ObjectHash> {
 	widget_manager: IWidgetManager;
 }
 
+type Experimental = {
+	invoke: <T>(
+		name: string,
+		msg?: any,
+		buffers?: DataView[],
+	) => Promise<[T, DataView[]]>;
+};
+
 export interface RenderProps<T extends ObjectHash = ObjectHash> {
 	model: AnyModel<T>;
 	el: HTMLElement;
+	experimental: Experimental;
 }
 
 export interface Render<T extends ObjectHash = ObjectHash> {
@@ -54,6 +63,7 @@ export interface Render<T extends ObjectHash = ObjectHash> {
 
 export interface InitializeProps<T extends ObjectHash = ObjectHash> {
 	model: AnyModel<T>;
+	experimental: Experimental;
 }
 
 export interface Initialize<T extends ObjectHash = ObjectHash> {