@marimo-team/islands 0.18.2 → 0.18.4

package/src/utils/data-views.ts

@@ -1,19 +1,131 @@
 /* Copyright 2024 Marimo. All rights reserved. */
-import { set } from "lodash-es";
+import { get, set } from "lodash-es";
 import { invariant } from "./invariant";
+import {
+  type Base64String,
+  binaryToByteString,
+  byteStringToBinary,
+  typedAtob,
+  typedBtoa,
+} from "./json/base64";
 import { Logger } from "./Logger";
 
 /**
- * Update the object with DataView buffers at the specified paths.
+ * Convert a DataView to a base64 string.
  */
-export function updateBufferPaths<T extends Record<string, unknown>>(
+export function dataViewToBase64(dataView: DataView): Base64String {
+  const bytes = new Uint8Array(
+    dataView.buffer,
+    dataView.byteOffset,
+    dataView.byteLength,
+  );
+  const byteString = binaryToByteString(bytes);
+  return typedBtoa(byteString);
+}
+
+/**
+ * Recursively find all DataViews in an object and return their paths.
+ *
+ * This mirrors ipywidgets' _separate_buffers logic.
+ */
+function findDataViewPaths(
+  obj: unknown,
+  currentPath: (string | number)[] = [],
+): (string | number)[][] {
+  const paths: (string | number)[][] = [];
+
+  if (obj instanceof DataView) {
+    paths.push(currentPath);
+  } else if (Array.isArray(obj)) {
+    for (let i = 0; i < obj.length; i++) {
+      paths.push(...findDataViewPaths(obj[i], [...currentPath, i]));
+    }
+  } else if (obj !== null && typeof obj === "object") {
+    for (const [key, value] of Object.entries(obj)) {
+      paths.push(...findDataViewPaths(value, [...currentPath, key]));
+    }
+  }
+
+  return paths;
+}
+
+/**
+ * Serialize DataView buffers to base64 strings.
+ *
+ * Finds all DataViews in the object.
+ */
+export function serializeBuffersToBase64<T extends Record<string, unknown>>(
   inputObject: T,
-  bufferPaths: readonly (readonly (string | number)[])[],
-  buffers: readonly DataView[],
-): T {
-  // If no buffer paths, return the original object
+): WireFormat<T> {
+  // Dynamically find all DataView paths instead of using fixed bufferPaths
+  const dataViewPaths = findDataViewPaths(inputObject);
+
+  if (dataViewPaths.length === 0) {
+    return { state: inputObject, buffers: [], bufferPaths: [] };
+  }
+
+  const state = structuredClone(inputObject);
+  const buffers: Base64String[] = [];
+  const bufferPaths: (string | number)[][] = [];
+
+  for (const bufferPath of dataViewPaths) {
+    const dataView = get(inputObject, bufferPath);
+    if (dataView instanceof DataView) {
+      const base64 = dataViewToBase64(dataView);
+      buffers.push(base64);
+      bufferPaths.push(bufferPath);
+      set(state, bufferPath, base64);
+    }
+  }
+
+  return { state, buffers, bufferPaths };
+}
+
+/**
+ * Wire format for anywidget state with binary data.
+ * Buffers can be either base64 strings (from network) or DataViews (in-memory).
+ */
+export interface WireFormat<T = Record<string, unknown>> {
+  state: T;
+  bufferPaths: (string | number)[][];
+  buffers: Base64String[];
+}
+
+/**
+ * Check if an object is in wire format.
+ */
+export function isWireFormat<T = Record<string, unknown>>(
+  obj: unknown,
+): obj is WireFormat<T> {
+  return (
+    obj !== null &&
+    typeof obj === "object" &&
+    "state" in obj &&
+    "bufferPaths" in obj &&
+    "buffers" in obj
+  );
+}
+
+/**
+ * Decode wire format or insert DataViews at specified paths.
+ *
+ * Accepts either:
+ * 1. Wire format: { state, bufferPaths, buffers } where buffers are base64 strings
+ * 2. Direct format: { state, bufferPaths, buffers } where buffers are DataViews
+ *
+ * For ndarray-like structures {view: null, dtype, shape}, we insert the DataView
+ * at the 'view' key, preserving the structure for round-tripping.
+ */
+export function decodeFromWire<T extends Record<string, unknown>>(input: {
+  state: T;
+  bufferPaths?: (string | number)[][];
+  buffers?: readonly (DataView | Base64String)[];
+}): T {
+  const { state, bufferPaths, buffers } = input;
+
+  // If no buffer paths, return the original state
   if (!bufferPaths || bufferPaths.length === 0) {
-    return inputObject;
+    return state;
   }
 
   // If has buffers, assert they are the same size
@@ -24,18 +136,24 @@ export function updateBufferPaths<T extends Record<string, unknown>>(
     );
   }
 
-  let object = structuredClone(inputObject);
+  const out = structuredClone(state);
 
   for (const [i, bufferPath] of bufferPaths.entries()) {
-    // If buffers exists, we use that value
-    // Otherwise we grab it from inside the inputObject
-    const dataView = buffers[i];
-    if (!dataView) {
-      Logger.warn("Could not find buffer at path", bufferPath);
+    const buffer = buffers?.[i];
+
+    if (buffer == null) {
+      Logger.warn("[anywidget] Could not find buffer at path", bufferPath);
       continue;
     }
-    object = set(object, bufferPath, dataView);
+
+    // Handle both base64 strings (from wire format) and DataViews (direct usage)
+    if (typeof buffer === "string") {
+      const bytes = byteStringToBinary(typedAtob(buffer));
+      set(out, bufferPath, new DataView(bytes.buffer));
+    } else {
+      set(out, bufferPath, buffer);
+    }
   }
 
-  return object;
+  return out;
 }

package/src/utils/json/base64.ts

@@ -55,6 +55,14 @@ export function byteStringToBinary(bytes: ByteString): Uint8Array {
   return Uint8Array.from(bytes, (c) => c.charCodeAt(0));
 }
 
+export function binaryToByteString(binary: Uint8Array): ByteString {
+  let result = "";
+  for (const byte of binary) {
+    result += String.fromCharCode(byte);
+  }
+  return result as ByteString;
+}
+
 export function safeExtractSetUIElementMessageBuffers(
   op: OperationMessageData<"send-ui-element-message">,
 ): readonly DataView[] {