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

.changeset/nervous-goats-rescue.md

@@ -0,0 +1,27 @@
+---
+"anywidget": patch
+"@anywidget/types": patch
+---
+
+Add experimental reducer API for sending/await-ing custom messages
+
+Introduces a unified API for dispatching messages from the front end, and
+await-ing response from Python. This removes a lot of boilerplate required for
+this kind of pattern previously. This API is experimental and opt-in, only if
+`_experimental_anywidget_reducer` is implemented on the anywidget subclass.
+
+```py
+class Widget(anywidget.AnyWidget):
+    _esm = """
+    async function render({ model, el, experimental }) {
+      let [response, buffers] = await experimental.dispatch("ping");
+      // Handle the response
+      console.log(response) // pong
+    }
+    export default { render };
+    """
+
+    def _experimental_anywidget_reducer(self, action, buffers):
+        assert action == "ping"
+        return "pong", []
+```

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
 
@@ -60,3 +60,20 @@ class AnywidgetProtocol(Protocol):
     """Anywidget classes have a MimeBundleDescriptor at `_repr_mimebundle_`."""
 
     _repr_mimebundle_: MimeBundleDescriptor
+
+
+class AnywidgetReducerProtocol(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:
+        ...
+
+    def _experimental_anywidget_reducer(
+        self, action: str | dict | list, buffers: list[bytes]
+    ) -> tuple[str | dict | list, list[bytes]]:
+        ...

anywidget/experimental.py

@@ -8,6 +8,9 @@
 
 from ._descriptor import MimeBundleDescriptor
 
+if typing.TYPE_CHECKING:  # pragma: no cover
+    from ._protocols import AnywidgetReducerProtocol
+
 __all__ = ["dataclass", "widget", "MimeBundleDescriptor"]
 
 _T = typing.TypeVar("_T")
@@ -103,3 +106,31 @@ def _decorator(cls: T) -> T:
         return cls
 
     return _decorator(cls) if cls is not None else _decorator  # type: ignore
+
+
+def _register_experimental_custom_message_reducer(
+    widget: AnywidgetReducerProtocol,
+) -> None:
+    """Register a custom message reducer for a widget if it implements the protocol."""
+    # Only add the reducer if it doesn't already exist
+    if not hasattr(widget, "_experimental_anywidget_reducer"):
+        return
+
+    def handle_anywidget_dispatch(
+        self: AnywidgetReducerProtocol, msg: str | list | dict, buffers: list[bytes]
+    ) -> None:
+        if not isinstance(msg, dict) or msg.get("kind") != "anywidget-dispatch":
+            return
+        response, buffers = widget._experimental_anywidget_reducer(
+            msg["action"], buffers
+        )
+        self.send(
+            {
+                "id": msg["id"],
+                "kind": "anywidget-dispatch-response",
+                "response": response,
+            },
+            buffers,
+        )
+
+    widget.on_msg(handle_anywidget_dispatch)

anywidget/widget.py

@@ -17,6 +17,7 @@
     try_file_contents,
 )
 from ._version import __version__
+from .experimental import _register_experimental_custom_message_reducer
 
 
 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_experimental_custom_message_reducer(self)
 
     def __init_subclass__(cls, **kwargs: dict) -> None:
         """Coerces _esm and _css to FileContents if they are files."""

packages/anywidget/src/widget.js

@@ -242,6 +242,35 @@ function throw_anywidget_error(source) {
 	throw source;
 }
 
+/**
+ * @template T
+ * @param {import("@anywidget/types").AnyModel} model
+ * @param {any} [action]
+ * @param {{ timeout?: number }} [options]
+ */
+export function dispatch(model, action, { timeout = 3000 } = {}) {
+	let id = Date.now().toString(36);
+	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-dispatch-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-dispatch", action });
+	});
+}
+
 class Runtime {
 	/** @type {() => void} */
 	#disposer = () => {};
@@ -281,6 +310,9 @@ class Runtime {
 					let widget = await load_widget(update);
 					cleanup = await widget.initialize?.({
 						model: model_proxy(model, INITIALIZE_MARKER),
+						experimental: {
+							dispatch: dispatch.bind(null, model),
+						},
 					});
 					return ok(widget);
 				} catch (e) {
@@ -319,6 +351,9 @@ class Runtime {
 						cleanup = await widget.render?.({
 							model: model_proxy(model, view),
 							el: view.el,
+							experimental: {
+								dispatch: dispatch.bind(null, model),
+							},
 						});
 					} catch (e) {
 						throw_anywidget_error(e);

packages/types/index.ts

@@ -43,9 +43,14 @@ export interface AnyModel<T extends ObjectHash = ObjectHash> {
 	widget_manager: IWidgetManager;
 }
 
+type Experimental = {
+	dispatch: <T>(action: any) => Promise<T>;
+};
+
 export interface RenderProps<T extends ObjectHash = ObjectHash> {
 	model: AnyModel<T>;
 	el: HTMLElement;
+	experimental: Experimental;
 }
 
 export interface Render<T extends ObjectHash = ObjectHash> {
@@ -54,6 +59,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> {

tests/test_widget.py

@@ -267,3 +267,30 @@ class Widget(anywidget.AnyWidget):
         """
 
     assert Widget()._repr_mimebundle_() is None
+
+
+def test_anywidget_reducer_not_registered_by_default():
+    class Widget(anywidget.AnyWidget):
+        _esm = "export default { render({ model, el }) { el.innerText = 'Hello, world'; } }"
+
+    w = Widget()
+    assert len(w._msg_callbacks.callbacks) == 0
+
+
+def test_anywidget_reducer_registers_callback():
+    class Widget(anywidget.AnyWidget):
+        _esm = """
+        export default {
+            async render({ model, el, experimental }) {
+                let response = await experimental.dispatch("ping");
+                console.log(response); // pong
+            }
+        }
+        """
+
+        def _experimental_anywidget_reducer(self, action, buffers):
+            assert action == "ping"
+            return "pong", []
+
+    w = Widget()
+    assert len(w._msg_callbacks.callbacks) == 1