From 24afa7898491db33766ce712b63e5e54b714fe67 Mon Sep 17 00:00:00 2001 From: Trevor Manz Date: Wed, 24 Sep 2025 13:47:30 -0400 Subject: [PATCH 01/13] Add test for `SendUIElementMessage.broadcast()` payload --- tests/_messaging/test_ops.py | 30 ++++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/tests/_messaging/test_ops.py b/tests/_messaging/test_ops.py index 2efe2b276d4..096fb2f445b 100644 --- a/tests/_messaging/test_ops.py +++ b/tests/_messaging/test_ops.py @@ -7,6 +7,7 @@ from marimo._messaging.ops import ( CellOp, InstallingPackageAlert, + SendUIElementMessage, StartupLogs, VariableValue, ) @@ -99,3 +100,32 @@ def test_installing_package_alert_with_logs() -> None: assert alert.packages == packages assert alert.logs == logs assert alert.log_status == "start" + + +def test_send_ui_element_message_broadcast() -> None: + """Test SendUIElementMessage broadcasting and serialization.""" + stream = MockStream() + + msg = SendUIElementMessage( + ui_element="test_element", + model_id=None, + message={"action": "update", "value": 42}, + buffers=[b"buffer1", b"buffer2"], + ) + + msg.broadcast(stream=stream) + + assert len(stream.messages) == 1 + + assert stream.operations[0] == { + "op": "send-ui-element-message", + "ui_element": "test_element", + "model_id": None, + "message": {"action": "update", "value": 42}, + "buffers": [ + "YnVmZmVyMQ==", + "YnVmZmVyMg==", + ], + } + + assert stream.parsed_operations[0] == msg From 860837791510ffa2e45af6d613557d366fc7f887 Mon Sep 17 00:00:00 2001 From: Trevor Manz Date: Wed, 24 Sep 2025 13:47:30 -0400 Subject: [PATCH 02/13] Test snapshot of anwidget `data-initial-value` --- tests/_plugins/ui/_impl/test_anywidget.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/tests/_plugins/ui/_impl/test_anywidget.py b/tests/_plugins/ui/_impl/test_anywidget.py index de44c5f7fae..847a3bd3cc1 100644 --- a/tests/_plugins/ui/_impl/test_anywidget.py +++ b/tests/_plugins/ui/_impl/test_anywidget.py @@ -300,6 +300,13 @@ class BufferWidget(_anywidget.AnyWidget): assert wrapped._initial_value == {"array": data} assert wrapped._component_args["buffer-paths"] == [["array"]] + # test buffers are inlined as base64 inplace + assert ( + "data-initial-value='{"array":"AQIDBA=="}'" + in wrapped.text + ) + assert "data-buffer-paths='[["array"]]'" in wrapped.text + # Test updating the buffer new_data = bytes([5, 6, 7, 8]) wrapped.array = new_data From 96ab4b204f1f6cb1fb4fb41fd3c2784170655be9 Mon Sep 17 00:00:00 2001 From: Trevor Manz Date: Wed, 24 Sep 2025 13:47:30 -0400 Subject: [PATCH 03/13] Require strict `DataView` when decoding operation messages --- frontend/src/core/dom/events.ts | 2 +- frontend/src/core/dom/uiregistry.ts | 10 ++---- frontend/src/core/islands/main.ts | 8 +++-- .../src/core/websocket/useMarimoWebSocket.tsx | 7 +++-- .../impl/anywidget/AnyWidgetPlugin.tsx | 31 +++++++++++++------ frontend/src/plugins/impl/anywidget/model.ts | 5 ++- frontend/src/utils/data-views.ts | 25 ++++----------- frontend/src/utils/json/base64.ts | 13 ++++++++ 8 files changed, 57 insertions(+), 44 deletions(-) diff --git a/frontend/src/core/dom/events.ts b/frontend/src/core/dom/events.ts index 34d7529157e..83e94a115e8 100644 --- a/frontend/src/core/dom/events.ts +++ b/frontend/src/core/dom/events.ts @@ -43,7 +43,7 @@ export const MarimoIncomingMessageEvent = defineCustomEvent( )<{ objectId: UIElementId; message: unknown; - buffers: DataView[] | undefined; + buffers: readonly DataView[]; }>(); export type MarimoIncomingMessageEventType = ReturnType< typeof MarimoIncomingMessageEvent.create diff --git a/frontend/src/core/dom/uiregistry.ts b/frontend/src/core/dom/uiregistry.ts index 32ce056c382..39c3c9cc72c 100644 --- a/frontend/src/core/dom/uiregistry.ts +++ b/frontend/src/core/dom/uiregistry.ts @@ -1,8 +1,5 @@ /* Copyright 2024 Marimo. All rights reserved. */ -import { byteStringToDataView } from "@/utils/data-views"; -import type { Base64String } from "@/utils/json/base64"; -import { typedAtob } from "@/utils/json/base64"; import { Logger } from "@/utils/Logger"; import type { CellId, UIElementId } from "../cells/ids"; import { @@ -136,15 +133,12 @@ export class UIElementRegistry { broadcastMessage( objectId: UIElementId, message: unknown, - buffers: Base64String[] | undefined | null, + buffers: readonly DataView[], ): void { const entry = this.entries.get(objectId); if (entry === undefined) { Logger.warn("UIElementRegistry missing entry", objectId); } else { - const toDataView = (base64: Base64String) => { - return byteStringToDataView(typedAtob(base64)); - }; entry.elements.forEach((element) => { element.dispatchEvent( MarimoIncomingMessageEvent.create({ @@ -153,7 +147,7 @@ export class UIElementRegistry { detail: { objectId: objectId, message: message, - buffers: buffers ? buffers.map(toDataView) : undefined, + buffers: buffers, }, }), ); diff --git a/frontend/src/core/islands/main.ts b/frontend/src/core/islands/main.ts index 721913233ee..5de965485fa 100644 --- a/frontend/src/core/islands/main.ts +++ b/frontend/src/core/islands/main.ts @@ -16,7 +16,11 @@ import { renderHTML } from "@/plugins/core/RenderHTML"; import { initializePlugins } from "@/plugins/plugins"; import { logNever } from "@/utils/assertNever"; import { Functions } from "@/utils/functions"; -import type { Base64String } from "@/utils/json/base64"; +import { + type Base64String, + base64StringToDataView, + safeExtractSetUIElementMessageBuffers, +} from "@/utils/json/base64"; import { jsonParseWithSpecialChar } from "@/utils/json/json-parser"; import { Logger } from "@/utils/Logger"; import { @@ -145,7 +149,7 @@ export async function initialize() { UI_ELEMENT_REGISTRY.broadcastMessage( msg.data.ui_element as UIElementId, msg.data.message, - msg.data.buffers as Base64String[], + safeExtractSetUIElementMessageBuffers(msg.data), ); return; diff --git a/frontend/src/core/websocket/useMarimoWebSocket.tsx b/frontend/src/core/websocket/useMarimoWebSocket.tsx index e0cb61180ba..ee15c432f8f 100644 --- a/frontend/src/core/websocket/useMarimoWebSocket.tsx +++ b/frontend/src/core/websocket/useMarimoWebSocket.tsx @@ -16,7 +16,10 @@ import { } from "@/plugins/impl/anywidget/model"; import { logNever } from "@/utils/assertNever"; import { prettyError } from "@/utils/errors"; -import type { Base64String, JsonString } from "@/utils/json/base64"; +import { + type JsonString, + safeExtractSetUIElementMessageBuffers, +} from "@/utils/json/base64"; import { jsonParseWithSpecialChar } from "@/utils/json/json-parser"; import { Logger } from "@/utils/Logger"; import { reloadSafe } from "@/utils/reload-safe"; @@ -112,7 +115,7 @@ export function useMarimoWebSocket(opts: { const modelId = msg.data.model_id; const uiElement = msg.data.ui_element; const message = msg.data.message; - const buffers = (msg.data.buffers ?? []) as Base64String[]; + const buffers = safeExtractSetUIElementMessageBuffers(msg.data); if (modelId && isMessageWidgetState(message)) { handleWidgetMessage({ diff --git a/frontend/src/plugins/impl/anywidget/AnyWidgetPlugin.tsx b/frontend/src/plugins/impl/anywidget/AnyWidgetPlugin.tsx index 5c561ec2ca5..d6787f20c74 100644 --- a/frontend/src/plugins/impl/anywidget/AnyWidgetPlugin.tsx +++ b/frontend/src/plugins/impl/anywidget/AnyWidgetPlugin.tsx @@ -2,8 +2,8 @@ /* eslint-disable @typescript-eslint/no-explicit-any */ import type { AnyWidget, Experimental } from "@anywidget/types"; -import { isEqual } from "lodash-es"; -import { useEffect, useMemo, useRef } from "react"; +import { get, isEqual, set } from "lodash-es"; +import { useEffect, useRef } from "react"; import { z } from "zod"; import { MarimoIncomingMessageEvent } from "@/core/dom/events"; import { asRemoteURL } from "@/core/runtime/config"; @@ -16,7 +16,11 @@ import { import { createPlugin } from "@/plugins/core/builder"; import { rpc } from "@/plugins/core/rpc"; import type { IPluginProps } from "@/plugins/types"; -import { updateBufferPaths } from "@/utils/data-views"; +import { + type Base64String, + byteStringToBinary, + typedAtob, +} from "@/utils/json/base64"; import { Logger } from "@/utils/Logger"; import { ErrorBanner } from "../common/error-banner"; import { MODEL_MANAGER, Model } from "./model"; @@ -58,7 +62,7 @@ export const AnyWidgetPlugin = createPlugin("marimo-anywidget") type Props = IPluginProps; const AnyWidgetSlot = (props: Props) => { - const { css, jsUrl, jsHash, bufferPaths } = props.data; + const { css, jsUrl, jsHash, bufferPaths, initialValue } = props.data; // JS is an ESM file with a render function on it // export function render({ model, el }) { // ... @@ -85,10 +89,6 @@ const AnyWidgetSlot = (props: Props) => { } }, [hasError, jsUrl]); - const valueWithBuffer = useMemo(() => { - return updateBufferPaths(props.value, bufferPaths); - }, [props.value, bufferPaths]); - // Mount the CSS useEffect(() => { const shadowRoot = props.host.shadowRoot; @@ -157,7 +157,7 @@ const AnyWidgetSlot = (props: Props) => { key={key} {...props} widget={module.default} - value={valueWithBuffer} + value={resolveInitialValue(initialValue, bufferPaths ?? [])} /> ); }; @@ -284,3 +284,16 @@ export const visibleForTesting = { isAnyWidgetModule, getDirtyFields, }; + +function resolveInitialValue( + raw: Record, + bufferPaths: ReadonlyArray>, +) { + const out = structuredClone(raw); + for (const bufferPath of bufferPaths) { + const base64String: Base64String = get(raw, bufferPath); + const bytes = byteStringToBinary(typedAtob(base64String)); + set(out, bufferPath, new DataView(bytes.buffer)); + } + return out; +} diff --git a/frontend/src/plugins/impl/anywidget/model.ts b/frontend/src/plugins/impl/anywidget/model.ts index 9f913cea73d..2371d0251ba 100644 --- a/frontend/src/plugins/impl/anywidget/model.ts +++ b/frontend/src/plugins/impl/anywidget/model.ts @@ -10,7 +10,6 @@ import { assertNever } from "@/utils/assertNever"; import { Deferred } from "@/utils/Deferred"; import { updateBufferPaths } from "@/utils/data-views"; import { throwNotImplemented } from "@/utils/functions"; -import type { Base64String } from "@/utils/json/base64"; import { Logger } from "@/utils/Logger"; export type EventHandler = (...args: any[]) => void; @@ -171,7 +170,7 @@ export class Model> implements AnyModel { * When receiving a message from the backend. * We want to notify all listeners with `msg:custom` */ - receiveCustomMessage(message: any, buffers?: DataView[]): void { + receiveCustomMessage(message: any, buffers: readonly DataView[] = []): void { const response = AnyWidgetMessageSchema.safeParse(message); if (response.success) { const data = response.data; @@ -260,7 +259,7 @@ export async function handleWidgetMessage({ }: { modelId: string; msg: AnyWidgetMessage; - buffers: Base64String[]; + buffers: readonly DataView[]; modelManager: ModelManager; }): Promise { if (msg.method === "echo_update") { diff --git a/frontend/src/utils/data-views.ts b/frontend/src/utils/data-views.ts index 263ee4994f2..b2a958080df 100644 --- a/frontend/src/utils/data-views.ts +++ b/frontend/src/utils/data-views.ts @@ -1,16 +1,15 @@ /* Copyright 2024 Marimo. All rights reserved. */ -import { get, set } from "lodash-es"; +import { set } from "lodash-es"; import { invariant } from "./invariant"; -import { type Base64String, type ByteString, typedAtob } from "./json/base64"; import { Logger } from "./Logger"; /** - * Update the object with DataView buffers at the specified paths. + * Update the ob ect with DataView buffers at the specified paths. */ export function updateBufferPaths>( inputObject: T, bufferPaths: Array> | null | undefined, - buffers?: Base64String[] | null | undefined, + buffers: readonly DataView[], ): T { // If no buffer paths, return the original object if (!bufferPaths || bufferPaths.length === 0) { @@ -30,25 +29,13 @@ export function updateBufferPaths>( for (const [i, bufferPath] of bufferPaths.entries()) { // If buffers exists, we use that value // Otherwise we grab it from inside the inputObject - const bytes: ByteString = buffers - ? typedAtob(buffers[i]) - : get(object, bufferPath); - if (!bytes) { + const dataView = buffers[i]; + if (!dataView) { Logger.warn("Could not find buffer at path", bufferPath); continue; } - const buffer = byteStringToDataView(bytes); - object = set(object, bufferPath, buffer); + object = set(object, bufferPath, dataView); } return object; } - -export const byteStringToDataView = (bytes: ByteString) => { - const buffer = new ArrayBuffer(bytes.length); - const view = new DataView(buffer); - for (let i = 0; i < bytes.length; i++) { - view.setUint8(i, bytes.charCodeAt(i)); - } - return view; -}; diff --git a/frontend/src/utils/json/base64.ts b/frontend/src/utils/json/base64.ts index 05e76e6b49d..5d419abcd8d 100644 --- a/frontend/src/utils/json/base64.ts +++ b/frontend/src/utils/json/base64.ts @@ -1,4 +1,6 @@ /* Copyright 2024 Marimo. All rights reserved. */ + +import type { OperationMessageData } from "@/core/kernel/messages"; import type { TypedString } from "../typed"; export type JsonString = TypedString<"Json"> & { @@ -52,3 +54,14 @@ export function extractBase64FromDataURL(str: DataURLString): Base64String { export function byteStringToBinary(bytes: ByteString): Uint8Array { return Uint8Array.from(bytes, (c) => c.charCodeAt(0)); } + +export function safeExtractSetUIElementMessageBuffers( + op: OperationMessageData<"send-ui-element-message">, +): readonly DataView[] { + // @ts-expect-error - TypeScript doesn't know that these strings are actually base64 strings + const strs: Base64String[] = op.buffers ?? []; + return strs.map((str) => { + const bytes = byteStringToBinary(typedAtob(str)); + return new DataView(bytes.buffer); + }); +} From 0bd202876bc78e67e23425037bb28a0d79a219e2 Mon Sep 17 00:00:00 2001 From: Trevor Manz Date: Wed, 24 Sep 2025 13:47:30 -0400 Subject: [PATCH 04/13] Prefer `msgspec` base64 encoding for `SendUIElementMessage` ops --- marimo/_messaging/ops.py | 2 +- marimo/_plugins/ui/_core/ui_element.py | 5 +---- marimo/_plugins/ui/_impl/comm.py | 6 +----- 3 files changed, 3 insertions(+), 10 deletions(-) diff --git a/marimo/_messaging/ops.py b/marimo/_messaging/ops.py index 12f3c42af59..4da2e836f34 100644 --- a/marimo/_messaging/ops.py +++ b/marimo/_messaging/ops.py @@ -381,7 +381,7 @@ class SendUIElementMessage(Op, tag="send-ui-element-message"): ui_element: Optional[str] model_id: Optional[WidgetModelId] message: dict[str, Any] - buffers: Optional[list[str]] = None + buffers: Optional[list[bytes]] = None class Interrupted(Op, tag="interrupted"): diff --git a/marimo/_plugins/ui/_core/ui_element.py b/marimo/_plugins/ui/_core/ui_element.py index 35304f180c4..28456d97abf 100644 --- a/marimo/_plugins/ui/_core/ui_element.py +++ b/marimo/_plugins/ui/_core/ui_element.py @@ -2,7 +2,6 @@ from __future__ import annotations import abc -import base64 import copy import random import sys @@ -440,9 +439,7 @@ def _send_message( ui_element=self._id, model_id=None, message=message, - buffers=[ - base64.b64encode(buffer).decode() for buffer in (buffers or []) - ], + buffers=list(buffers or []), ).broadcast() def _update(self, value: S) -> None: diff --git a/marimo/_plugins/ui/_impl/comm.py b/marimo/_plugins/ui/_impl/comm.py index 127ef316279..baf74cf2fd7 100644 --- a/marimo/_plugins/ui/_impl/comm.py +++ b/marimo/_plugins/ui/_impl/comm.py @@ -1,7 +1,6 @@ # Copyright 2024 Marimo. All rights reserved. from __future__ import annotations -import base64 from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Callable, Optional, cast @@ -229,10 +228,7 @@ def flush(self) -> None: ui_element=self.ui_element_id, model_id=item.model_id, message=item.data, - buffers=[ - base64.b64encode(buffer).decode() - for buffer in item.buffers - ], + buffers=item.buffers, ).broadcast() # This is the method that ipywidgets.widgets.Widget uses to respond to From 8955765f8552cb50e2e72f8da1ed3b8200e87eb7 Mon Sep 17 00:00:00 2001 From: Trevor Manz Date: Wed, 24 Sep 2025 16:07:46 -0400 Subject: [PATCH 05/13] Add `uchimata` smoketest --- .../anywidget_smoke_tests/uchimata_example.py | 81 +++++++++++++++++++ 1 file changed, 81 insertions(+) create mode 100644 marimo/_smoke_tests/anywidget_smoke_tests/uchimata_example.py diff --git a/marimo/_smoke_tests/anywidget_smoke_tests/uchimata_example.py b/marimo/_smoke_tests/anywidget_smoke_tests/uchimata_example.py new file mode 100644 index 00000000000..734baf7389d --- /dev/null +++ b/marimo/_smoke_tests/anywidget_smoke_tests/uchimata_example.py @@ -0,0 +1,81 @@ +# /// script +# requires-python = ">=3.12" +# dependencies = [ +# "anywidget==0.9.18", +# "numpy==2.3.3", +# "polars==1.33.1", +# "traitlets==5.14.3", +# "uchimata==0.3.0", +# ] +# /// + +import marimo + +__generated_with = "0.16.2" +app = marimo.App(width="medium") + + +@app.cell +def _(): + import anywidget + import traitlets + + class Widget(anywidget.AnyWidget): + _esm = """ + export default { + render({ model, el }) { + const dataView = model.get("data"); + const bytes = new Uint8Array(dataView.buffer) + const decoded = new TextDecoder().decode(bytes); + el.innerText = decoded; + } + } + """ + data = traitlets.Any().tag(sync=True) + + # Should display "hello" + Widget(data=b"hello") + return + + +@app.cell +def _(): + import uchimata as uchi + import numpy as np + + BINS_NUM = 1000 + + # Step 1: Generate random structure, returns a 2D numpy array: + def make_random_3D_chromatin_structure(n): + position = np.array([0.0, 0.0, 0.0]) + positions = [position.copy()] + for _ in range(n): + step = np.random.choice( + [-1.0, 0.0, 1.0], size=3 + ) # Randomly choose to move left, right, up, down, forward, or backward + position += step + positions.append(position.copy()) + return np.array(positions) + + random_structure = make_random_3D_chromatin_structure(BINS_NUM) + + # Step 2: Display the structure in an uchimata widget + numbers = list(range(0, BINS_NUM + 1)) + vc = { + "color": { + "values": numbers, + "min": 0, + "max": BINS_NUM, + "colorScale": "Spectral", + }, + "scale": 0.01, + "links": True, + "mark": "sphere", + } + + uchi.Widget(random_structure, vc) + return + + +if __name__ == "__main__": + app.run() From 419c1862d7b080ec868ba18147aef5af86a2483f Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 24 Sep 2025 20:25:17 +0000 Subject: [PATCH 06/13] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- frontend/src/utils/data-views.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/frontend/src/utils/data-views.ts b/frontend/src/utils/data-views.ts index b2a958080df..07a673eaeec 100644 --- a/frontend/src/utils/data-views.ts +++ b/frontend/src/utils/data-views.ts @@ -4,7 +4,7 @@ import { invariant } from "./invariant"; import { Logger } from "./Logger"; /** - * Update the ob ect with DataView buffers at the specified paths. + * Update the ob etc with DataView buffers at the specified paths. */ export function updateBufferPaths>( inputObject: T, From 51d846d2307aadc1e06fa64e524e33d500dda1f1 Mon Sep 17 00:00:00 2001 From: Trevor Manz Date: Wed, 24 Sep 2025 16:38:14 -0400 Subject: [PATCH 07/13] Update frontend tests --- frontend/src/core/islands/main.ts | 6 +- .../__tests__/AnyWidgetPlugin.test.tsx | 2 +- .../impl/anywidget/__tests__/model.test.ts | 3 +- .../src/utils/__tests__/data-views.test.ts | 78 +++---------------- frontend/src/utils/data-views.ts | 2 +- 5 files changed, 14 insertions(+), 77 deletions(-) diff --git a/frontend/src/core/islands/main.ts b/frontend/src/core/islands/main.ts index 5de965485fa..288ca9ad2c0 100644 --- a/frontend/src/core/islands/main.ts +++ b/frontend/src/core/islands/main.ts @@ -16,11 +16,7 @@ import { renderHTML } from "@/plugins/core/RenderHTML"; import { initializePlugins } from "@/plugins/plugins"; import { logNever } from "@/utils/assertNever"; import { Functions } from "@/utils/functions"; -import { - type Base64String, - base64StringToDataView, - safeExtractSetUIElementMessageBuffers, -} from "@/utils/json/base64"; +import { safeExtractSetUIElementMessageBuffers } from "@/utils/json/base64"; import { jsonParseWithSpecialChar } from "@/utils/json/json-parser"; import { Logger } from "@/utils/Logger"; import { diff --git a/frontend/src/plugins/impl/anywidget/__tests__/AnyWidgetPlugin.test.tsx b/frontend/src/plugins/impl/anywidget/__tests__/AnyWidgetPlugin.test.tsx index 8beff09a51b..e20523bebcd 100644 --- a/frontend/src/plugins/impl/anywidget/__tests__/AnyWidgetPlugin.test.tsx +++ b/frontend/src/plugins/impl/anywidget/__tests__/AnyWidgetPlugin.test.tsx @@ -129,7 +129,7 @@ describe("LoadedSlot", () => { method: "update", state: { count: 10 }, }, - buffers: undefined, + buffers: [], }, bubbles: false, composed: true, diff --git a/frontend/src/plugins/impl/anywidget/__tests__/model.test.ts b/frontend/src/plugins/impl/anywidget/__tests__/model.test.ts index 0a6cd48b60e..bdb3f5f8699 100644 --- a/frontend/src/plugins/impl/anywidget/__tests__/model.test.ts +++ b/frontend/src/plugins/impl/anywidget/__tests__/model.test.ts @@ -9,7 +9,6 @@ import { vi, } from "vitest"; import { TestUtils } from "@/__tests__/test-helpers"; -import type { Base64String } from "@/utils/json/base64"; import { type AnyWidgetMessage, handleWidgetMessage, @@ -286,7 +285,7 @@ describe("ModelManager", () => { }: { modelId: string; message: AnyWidgetMessage; - buffers: Base64String[]; + buffers: readonly DataView[]; }) => { return handleWidgetMessage({ modelId, diff --git a/frontend/src/utils/__tests__/data-views.test.ts b/frontend/src/utils/__tests__/data-views.test.ts index f811cabd492..f1cc37da811 100644 --- a/frontend/src/utils/__tests__/data-views.test.ts +++ b/frontend/src/utils/__tests__/data-views.test.ts @@ -1,39 +1,11 @@ /* Copyright 2024 Marimo. All rights reserved. */ import { describe, expect, it } from "vitest"; -import { byteStringToDataView, updateBufferPaths } from "../data-views"; -import type { Base64String, ByteString } from "../json/base64"; +import { updateBufferPaths } from "../data-views"; describe("updateBufferPaths", () => { - it("should return the original object if bufferPaths is null", () => { + it("should return the original object if bufferPaths.length === 0", () => { const input = { a: 1, b: 2 }; - const result = updateBufferPaths(input, null); - expect(result).toEqual(input); - }); - - it("should update buffer paths correctly", () => { - const input = { - a: 1, - b: { - c: "Hello", - d: "World", - }, - }; - const bufferPaths = [ - ["b", "c"], - ["b", "d"], - ]; - const result = updateBufferPaths(input, bufferPaths); - - expect(result.a).toBe(1); - expect(result.b.c).toBeInstanceOf(DataView); - expect(result.b.d).toBeInstanceOf(DataView); - }); - - it("should handle non-existent paths", () => { - const input = { a: 1 }; - const bufferPaths = [["b", "c"]]; - const result = updateBufferPaths(input, bufferPaths); - + const result = updateBufferPaths(input, [], []); expect(result).toEqual(input); }); @@ -49,19 +21,12 @@ describe("updateBufferPaths", () => { ["b", "c"], ["b", "d"], ]; - const buffers: Base64String[] = [ - "SGVsbG8=" as Base64String, - "V29ybGQ=" as Base64String, - ]; // Base64 encoded "Hello" and "World" + const buffers = [ + new TextEncoder().encode("Hello"), + new TextEncoder().encode("World"), + ].map((b) => new DataView(b.buffer)); const result = updateBufferPaths(input, bufferPaths, buffers); - - expect(result.a).toBe(1); - const cView = result.b.c as unknown as DataView; - const dView = result.b.d as unknown as DataView; - expect(cView).toBeInstanceOf(DataView); - expect(dView).toBeInstanceOf(DataView); - expect(cView.byteLength).toBe(5); - expect(dView.byteLength).toBe(5); + expect(result).toMatchInlineSnapshot(); }); it("should throw error when buffers and paths length mismatch", () => { @@ -70,7 +35,7 @@ describe("updateBufferPaths", () => { ["b", "c"], ["b", "d"], ]; - const buffers: Base64String[] = ["SGVsbG8=" as Base64String]; // Only one buffer for two paths + const buffers = [new DataView(new ArrayBuffer())]; // Only one buffer for two paths expect(() => updateBufferPaths(input, bufferPaths, buffers)).toThrow( "Buffers and buffer paths not the same length", @@ -80,33 +45,10 @@ describe("updateBufferPaths", () => { it("should handle empty buffers array", () => { const input = { a: 1 }; const bufferPaths = [["b", "c"]]; - const buffers: Base64String[] = []; + const buffers: DataView[] = []; expect(() => updateBufferPaths(input, bufferPaths, buffers)).toThrow( "Buffers and buffer paths not the same length", ); }); }); - -describe("byteStringToDataView", () => { - it("should convert a base64 string to a DataView", () => { - const input = "Hello" as ByteString; - const result = byteStringToDataView(input); - - expect(result).toBeInstanceOf(DataView); - expect(result.byteLength).toBe(5); - expect(result.getUint8(0)).toBe(72); // 'H' - expect(result.getUint8(1)).toBe(101); // 'e' - expect(result.getUint8(2)).toBe(108); // 'l' - expect(result.getUint8(3)).toBe(108); // 'l' - expect(result.getUint8(4)).toBe(111); // 'o' - }); - - it("should handle empty string", () => { - const input = "" as ByteString; - const result = byteStringToDataView(input); - - expect(result).toBeInstanceOf(DataView); - expect(result.byteLength).toBe(0); - }); -}); diff --git a/frontend/src/utils/data-views.ts b/frontend/src/utils/data-views.ts index 07a673eaeec..197c57f3624 100644 --- a/frontend/src/utils/data-views.ts +++ b/frontend/src/utils/data-views.ts @@ -8,7 +8,7 @@ import { Logger } from "./Logger"; */ export function updateBufferPaths>( inputObject: T, - bufferPaths: Array> | null | undefined, + bufferPaths: ReadonlyArray>, buffers: readonly DataView[], ): T { // If no buffer paths, return the original object From f3701fd6fd11bc660c038aa458d71b4c3b5ca251 Mon Sep 17 00:00:00 2001 From: Trevor Manz Date: Wed, 24 Sep 2025 16:47:14 -0400 Subject: [PATCH 08/13] `useMemo` for initialValue --- .../impl/anywidget/AnyWidgetPlugin.tsx | 13 +++-- .../__tests__/AnyWidgetPlugin.test.tsx | 58 ++++++++++++++++++- 2 files changed, 66 insertions(+), 5 deletions(-) diff --git a/frontend/src/plugins/impl/anywidget/AnyWidgetPlugin.tsx b/frontend/src/plugins/impl/anywidget/AnyWidgetPlugin.tsx index d6787f20c74..9b305b41e72 100644 --- a/frontend/src/plugins/impl/anywidget/AnyWidgetPlugin.tsx +++ b/frontend/src/plugins/impl/anywidget/AnyWidgetPlugin.tsx @@ -3,7 +3,7 @@ import type { AnyWidget, Experimental } from "@anywidget/types"; import { get, isEqual, set } from "lodash-es"; -import { useEffect, useRef } from "react"; +import { useEffect, useRef, useMemo } from "react"; import { z } from "zod"; import { MarimoIncomingMessageEvent } from "@/core/dom/events"; import { asRemoteURL } from "@/core/runtime/config"; @@ -62,7 +62,12 @@ export const AnyWidgetPlugin = createPlugin("marimo-anywidget") type Props = IPluginProps; const AnyWidgetSlot = (props: Props) => { - const { css, jsUrl, jsHash, bufferPaths, initialValue } = props.data; + const { css, jsUrl, jsHash, bufferPaths } = props.data; + + const valueWithBuffers = useMemo(() => { + return resolveInitialValue(props.value, bufferPaths ?? []); + }, [props.value, bufferPaths]); + // JS is an ESM file with a render function on it // export function render({ model, el }) { // ... @@ -157,7 +162,7 @@ const AnyWidgetSlot = (props: Props) => { key={key} {...props} widget={module.default} - value={resolveInitialValue(initialValue, bufferPaths ?? [])} + value={valueWithBuffers} /> ); }; @@ -285,7 +290,7 @@ export const visibleForTesting = { getDirtyFields, }; -function resolveInitialValue( +export function resolveInitialValue( raw: Record, bufferPaths: ReadonlyArray>, ) { diff --git a/frontend/src/plugins/impl/anywidget/__tests__/AnyWidgetPlugin.test.tsx b/frontend/src/plugins/impl/anywidget/__tests__/AnyWidgetPlugin.test.tsx index e20523bebcd..05e720b1411 100644 --- a/frontend/src/plugins/impl/anywidget/__tests__/AnyWidgetPlugin.test.tsx +++ b/frontend/src/plugins/impl/anywidget/__tests__/AnyWidgetPlugin.test.tsx @@ -5,7 +5,11 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; import { TestUtils } from "@/__tests__/test-helpers"; import type { UIElementId } from "@/core/cells/ids"; import { MarimoIncomingMessageEvent } from "@/core/dom/events"; -import { getDirtyFields, visibleForTesting } from "../AnyWidgetPlugin"; +import { + getDirtyFields, + visibleForTesting, + resolveInitialValue, +} from "../AnyWidgetPlugin"; import { Model } from "../model"; const { LoadedSlot } = visibleForTesting; @@ -179,3 +183,55 @@ describe("LoadedSlot", () => { }); }); }); + +describe("resolveInitialValue", () => { + it("should convert base64 strings to DataView at specified paths", () => { + const result = resolveInitialValue( + { + a: 10, + b: "aGVsbG8=", // "hello" in base64 + c: [1, "d29ybGQ="], // "world" in base64 + d: { + foo: "bWFyaW1vCg==", // "marimo" in base64 + baz: 20, + }, + }, + [["b"], ["c", 1], ["d", "foo"]], + ); + + expect(result).toMatchInlineSnapshot(` + { + "a": 10, + "b": DataView [ + 104, + 101, + 108, + 108, + 111, + ], + "c": [ + 1, + DataView [ + 119, + 111, + 114, + 108, + 100, + ], + ], + "d": { + "baz": 20, + "foo": DataView [ + 109, + 97, + 114, + 105, + 109, + 111, + 10, + ], + }, + } + `); + }); +}); From 18b78bfbc17a49385e5e1e8d6e331defd263def5 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 24 Sep 2025 20:51:29 +0000 Subject: [PATCH 09/13] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../impl/anywidget/AnyWidgetPlugin.tsx | 2 +- .../__tests__/AnyWidgetPlugin.test.tsx | 2 +- .../src/utils/__tests__/data-views.test.ts | 22 ++++++++++++++++- .../anywidget_smoke_tests/mosaic_example.py | 24 ++++++++++++++++++- 4 files changed, 46 insertions(+), 4 deletions(-) diff --git a/frontend/src/plugins/impl/anywidget/AnyWidgetPlugin.tsx b/frontend/src/plugins/impl/anywidget/AnyWidgetPlugin.tsx index 9b305b41e72..e073a95ed1b 100644 --- a/frontend/src/plugins/impl/anywidget/AnyWidgetPlugin.tsx +++ b/frontend/src/plugins/impl/anywidget/AnyWidgetPlugin.tsx @@ -3,7 +3,7 @@ import type { AnyWidget, Experimental } from "@anywidget/types"; import { get, isEqual, set } from "lodash-es"; -import { useEffect, useRef, useMemo } from "react"; +import { useEffect, useMemo, useRef } from "react"; import { z } from "zod"; import { MarimoIncomingMessageEvent } from "@/core/dom/events"; import { asRemoteURL } from "@/core/runtime/config"; diff --git a/frontend/src/plugins/impl/anywidget/__tests__/AnyWidgetPlugin.test.tsx b/frontend/src/plugins/impl/anywidget/__tests__/AnyWidgetPlugin.test.tsx index 05e720b1411..738c70046b1 100644 --- a/frontend/src/plugins/impl/anywidget/__tests__/AnyWidgetPlugin.test.tsx +++ b/frontend/src/plugins/impl/anywidget/__tests__/AnyWidgetPlugin.test.tsx @@ -7,8 +7,8 @@ import type { UIElementId } from "@/core/cells/ids"; import { MarimoIncomingMessageEvent } from "@/core/dom/events"; import { getDirtyFields, - visibleForTesting, resolveInitialValue, + visibleForTesting, } from "../AnyWidgetPlugin"; import { Model } from "../model"; diff --git a/frontend/src/utils/__tests__/data-views.test.ts b/frontend/src/utils/__tests__/data-views.test.ts index f1cc37da811..cfae0f8c680 100644 --- a/frontend/src/utils/__tests__/data-views.test.ts +++ b/frontend/src/utils/__tests__/data-views.test.ts @@ -26,7 +26,27 @@ describe("updateBufferPaths", () => { new TextEncoder().encode("World"), ].map((b) => new DataView(b.buffer)); const result = updateBufferPaths(input, bufferPaths, buffers); - expect(result).toMatchInlineSnapshot(); + expect(result).toMatchInlineSnapshot(` + { + "a": 1, + "b": { + "c": DataView [ + 72, + 101, + 108, + 108, + 111, + ], + "d": DataView [ + 87, + 111, + 114, + 108, + 100, + ], + }, + } + `); }); it("should throw error when buffers and paths length mismatch", () => { diff --git a/marimo/_smoke_tests/anywidget_smoke_tests/mosaic_example.py b/marimo/_smoke_tests/anywidget_smoke_tests/mosaic_example.py index c28b98cec5e..44e43c78a79 100644 --- a/marimo/_smoke_tests/anywidget_smoke_tests/mosaic_example.py +++ b/marimo/_smoke_tests/anywidget_smoke_tests/mosaic_example.py @@ -5,13 +5,15 @@ # "mosaic-widget", # "marimo", # "pyyaml", +# "quak==0.3.2", +# "polars==1.33.1", # ] # /// # Copyright 2024 Marimo. All rights reserved. import marimo -__generated_with = "0.15.5" +__generated_with = "0.16.2" app = marimo.App(width="medium") @@ -53,5 +55,25 @@ def _(w): return +@app.cell +def _(): + import quak + return (quak,) + + +@app.cell +def _(quak): + import polars as pl + + _df = pl.read_parquet("https://github.com/uwdata/mosaic/raw/main/data/athletes.parquet") + quak.Widget(_df) + return + + +@app.cell +def _(): + return + + if __name__ == "__main__": app.run() From 4d7332b5612afe8ef4d3f0f48cff04574590c558 Mon Sep 17 00:00:00 2001 From: Trevor Manz Date: Wed, 24 Sep 2025 17:01:35 -0400 Subject: [PATCH 10/13] Fix tests with explicit buffers --- frontend/src/plugins/impl/anywidget/__tests__/model.test.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/frontend/src/plugins/impl/anywidget/__tests__/model.test.ts b/frontend/src/plugins/impl/anywidget/__tests__/model.test.ts index bdb3f5f8699..cc3265068f1 100644 --- a/frontend/src/plugins/impl/anywidget/__tests__/model.test.ts +++ b/frontend/src/plugins/impl/anywidget/__tests__/model.test.ts @@ -245,7 +245,7 @@ describe("Model", () => { content, }); - expect(callback).toHaveBeenCalledWith(content, undefined); + expect(callback).toHaveBeenCalledWith(content, []); }); it("should handle custom messages with buffers", () => { @@ -352,7 +352,7 @@ describe("ModelManager", () => { message: { method: "custom", content: { count: 1 } }, buffers: [], }); - expect(callback).toHaveBeenCalledWith({ count: 1 }, undefined); + expect(callback).toHaveBeenCalledWith({ count: 1 }, []); }); it("should handle close messages", async () => { From 9a00627f6e0e7e0d401fb3796237f91ac2bc75dc Mon Sep 17 00:00:00 2001 From: Trevor Manz Date: Wed, 24 Sep 2025 19:32:13 -0400 Subject: [PATCH 11/13] Update frontend/src/utils/data-views.ts Co-authored-by: Myles Scolnick --- frontend/src/utils/data-views.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/frontend/src/utils/data-views.ts b/frontend/src/utils/data-views.ts index 197c57f3624..fca9b52834d 100644 --- a/frontend/src/utils/data-views.ts +++ b/frontend/src/utils/data-views.ts @@ -4,7 +4,7 @@ import { invariant } from "./invariant"; import { Logger } from "./Logger"; /** - * Update the ob etc with DataView buffers at the specified paths. + * Update the object with DataView buffers at the specified paths. */ export function updateBufferPaths>( inputObject: T, From 6a303809a1417dfdb113dfb0b2eb725e27de0cf5 Mon Sep 17 00:00:00 2001 From: Trevor Manz Date: Wed, 24 Sep 2025 19:36:57 -0400 Subject: [PATCH 12/13] Update `api.yaml` --- packages/openapi/api.yaml | 1 + 1 file changed, 1 insertion(+) diff --git a/packages/openapi/api.yaml b/packages/openapi/api.yaml index 4201490617f..1ae97b810b9 100644 --- a/packages/openapi/api.yaml +++ b/packages/openapi/api.yaml @@ -2949,6 +2949,7 @@ components: buffers: anyOf: - items: + contentEncoding: base64 type: string type: array - type: 'null' From 352a27d3ab74a5cb9b3a1b630b645e5d28b78d6d Mon Sep 17 00:00:00 2001 From: Trevor Manz Date: Wed, 24 Sep 2025 20:19:53 -0400 Subject: [PATCH 13/13] Update python test with correct `SendUIElementMessage` checks --- tests/_plugins/ui/_impl/test_comm.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/tests/_plugins/ui/_impl/test_comm.py b/tests/_plugins/ui/_impl/test_comm.py index 3bbb9648a9e..8e639781483 100644 --- a/tests/_plugins/ui/_impl/test_comm.py +++ b/tests/_plugins/ui/_impl/test_comm.py @@ -163,13 +163,12 @@ def test_comm_flush(comm: MarimoComm): ) ) - with patch("marimo._messaging.ops.SendUIElementMessage") as mock_send: + with patch( + "marimo._messaging.ops.SendUIElementMessage" + ) as MockSendUIElementMessage: comm.flush() - mock_send.assert_called_once() - call_args = mock_send.call_args[1] + MockSendUIElementMessage.assert_called_once() + call_args = MockSendUIElementMessage.call_args[1] assert call_args["model_id"] == comm.comm_id assert call_args["message"] == test_data - assert len(call_args["buffers"]) == 1 - assert ( - call_args["buffers"][0] == "dGVzdF9idWZmZXI=" - ) # base64 of b"test_buffer" + assert call_args["buffers"] == [b"test_buffer"]