feat(experimental): Add Jupyter Widget (#176)

Author: manzt

.gitignore

@@ -29,3 +29,7 @@ example/*.zarr
 example/.ipynb_checkpoints/*
 example/data/**
 __pycache__
+
+.venv
+.ipynb_checkpoints
+dist/

python/README.md

@@ -0,0 +1,14 @@
+# vizarr
+
+```sh
+pip install vizarr
+```
+
+```python
+import vizarr
+import zarr
+
+viewer = vizarr.Viewer()
+viewer.add_image(source=zarr.open("path/to/ome.zarr"))
+viewer
+```

python/deno.json

@@ -0,0 +1,22 @@
+{
+	"lock": false,
+	"compilerOptions": {
+		"checkJs": true,
+		"allowJs": true,
+		"lib": [
+			"ES2022",
+			"DOM",
+			"DOM.Iterable"
+		]
+	},
+	"fmt": {
+		"useTabs": true
+	},
+	"lint": {
+		"rules": {
+			"exclude": [
+				"prefer-const"
+			]
+		}
+	}
+}

python/pyproject.toml

@@ -0,0 +1,15 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "vizarr"
+version = "0.0.1"
+dependencies = ["anywidget", "zarr"]
+
+[project.optional-dependencies]
+dev = ["watchfiles", "jupyterlab"]
+
+# automatically add the dev feature to the default env (e.g., hatch shell)
+[tool.hatch.envs.default]
+features = ["dev"]

python/src/vizarr/__init__.py

@@ -0,0 +1,10 @@
+import importlib.metadata
+
+try:
+    __version__ = importlib.metadata.version("vizarr")
+except importlib.metadata.PackageNotFoundError:
+    __version__ = "unknown"
+
+del importlib
+
+from ._widget import Viewer

python/src/vizarr/_widget.js

@@ -0,0 +1,123 @@
+import * as vizarr from "https://hms-dbmi.github.io/vizarr/index.js";
+import debounce from "https://esm.sh/just-debounce-it@3";
+
+/**
+ * @template T
+ * @param {import("npm:@anywidget/types").AnyModel} model
+ * @param {any} payload
+ * @param {{ timeout?: number }} [options]
+ * @returns {Promise<{ data: T, buffers: DataView[] }>}
+ */
+function send(model, payload, { timeout = 3000 } = {}) {
+	let uuid = globalThis.crypto.randomUUID();
+	return new Promise((resolve, reject) => {
+		let timer = setTimeout(() => {
+			reject(new Error(`Promise timed out after ${timeout} ms`));
+			model.off("msg:custom", handler);
+		}, timeout);
+		/**
+		 * @param {{ uuid: string, payload: T }} msg
+		 * @param {DataView[]} buffers
+		 */
+		function handler(msg, buffers) {
+			if (!(msg.uuid === uuid)) return;
+			clearTimeout(timer);
+			resolve({ data: msg.payload, buffers });
+			model.off("msg:custom", handler);
+		}
+		model.on("msg:custom", handler);
+		model.send({ payload, uuid });
+	});
+}
+
+/**
+ * @param {import("npm:@anywidget/types").AnyModel} model
+ * @param {string | { id: string }} source
+ */
+function get_source(model, source) {
+	if (typeof source === "string") {
+		return source;
+	}
+	// create a python
+	return {
+		/**
+		 * @param {string} key
+		 * @return {Promise<ArrayBuffer>}
+		 */
+		async getItem(key) {
+			const { data, buffers } = await send(model, {
+				type: "get",
+				source_id: source.id,
+				key,
+			});
+			if (!data.success) {
+				throw { __zarr__: "KeyError" };
+			}
+			return buffers[0].buffer;
+		},
+		/**
+		 * @param {string} key
+		 * @return {Promise<boolean>}
+		 */
+		async containsItem(key) {
+			const { data } = await send(model, {
+				type: "has",
+				source_id: source.id,
+				key,
+			});
+			return data;
+		},
+	};
+}
+
+/**
+ * @typedef Model
+ * @property {string} height
+ * @property {ViewState=} view_state
+ * @property {{ source: string | { id: string }}[]} _configs
+ */
+
+/**
+ * @typedef ViewState
+ * @property {number} zoom
+ * @property {[x: number, y: number]} target
+ */
+
+/** @type {import("npm:@anywidget/types").Render<Model>} */
+export function render({ model, el }) {
+	let div = document.createElement("div");
+	{
+		div.style.height = model.get("height");
+		div.style.backgroundColor = "black";
+		model.on("change:height", () => {
+			div.style.height = model.get("height");
+		});
+	}
+	let viewer = vizarr.createViewer(div);
+	{
+		model.on("change:view_state", () => {
+			viewer.setViewState(model.get("view_state"));
+		});
+		viewer.on(
+			"viewStateChange",
+			debounce((/** @type {ViewState} */ update) => {
+				model.set("view_state", update);
+				model.save_changes();
+			}, 200),
+		);
+	}
+	{
+		// sources are append-only now
+		for (const config of model.get("_configs")) {
+			const source = get_source(model, config.source);
+			viewer.addImage({ ...config, source });
+		}
+		model.on("change:_configs", () => {
+			const last = model.get("_configs").at(-1);
+			if (!last) return;
+			const source = get_source(model, last.source);
+			viewer.addImage({ ...last, source });
+		});
+	}
+	el.appendChild(div);
+}

python/src/vizarr/_widget.py

@@ -0,0 +1,67 @@
+import anywidget
+import traitlets
+import pathlib
+
+import zarr
+import numpy as np
+
+__all__ = ["Viewer"]
+
+
+def _store_keyprefix(obj):
+    # Just grab the store and key_prefix from zarr.Array and zarr.Group objects
+    if isinstance(obj, (zarr.Array, zarr.Group)):
+        return obj.store, obj._key_prefix
+
+    if isinstance(obj, np.ndarray):
+        # Create an in-memory store, and write array as as single chunk
+        store = {}
+        arr = zarr.create(
+            store=store, shape=obj.shape, chunks=obj.shape, dtype=obj.dtype
+        )
+        arr[:] = obj
+        return store, ""
+
+    if hasattr(obj, "__getitem__") and hasattr(obj, "__contains__"):
+        return obj, ""
+
+    raise TypeError("Cannot normalize store path")
+
+
+class Viewer(anywidget.AnyWidget):
+    _esm = pathlib.Path(__file__).parent / "_widget.js"
+    _configs = traitlets.List().tag(sync=True)
+    view_state = traitlets.Dict().tag(sync=True)
+    height = traitlets.Unicode("500px").tag(sync=True)
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self._store_paths = []
+        self.on_msg(self._handle_custom_msg)
+
+    def _handle_custom_msg(self, msg, buffers):
+        store, key_prefix = self._store_paths[msg["payload"]["source_id"]]
+        key = key_prefix + msg["payload"]["key"].lstrip("/")
+
+        if msg["payload"]["type"] == "has":
+            self.send({"uuid": msg["uuid"], "payload": key in store})
+            return
+
+        if msg["payload"]["type"] == "get":
+            try:
+                buffers = [store[key]]
+            except KeyError:
+                buffers = []
+            self.send(
+                {"uuid": msg["uuid"], "payload": {"success": len(buffers) == 1}},
+                buffers,
+            )
+            return
+
+    def add_image(self, source, **config):
+        if not isinstance(source, str):
+            store, key_prefix = _store_keyprefix(source)
+            source = {"id": len(self._store_paths)}
+            self._store_paths.append((store, key_prefix))
+        config["source"] = source
+        self._configs = self._configs + [config]