@marimo-team/frontend 0.14.18-dev5 → 0.14.18-dev51

package/src/core/dom/ui-element.ts

@@ -0,0 +1,249 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+import { isCustomMarimoElement } from "@/plugins/core/registerReactComponent";
+import { Functions } from "../../utils/functions";
+import { Logger } from "../../utils/Logger";
+import { UIElementId } from "../cells/ids";
+import { defineCustomElement } from "./defineCustomElement";
+import {
+  MarimoValueInputEvent,
+  type MarimoValueInputEventType,
+} from "./events";
+import { UI_ELEMENT_REGISTRY } from "./uiregistry";
+
+import "./ui-element.css";
+
+const UI_ELEMENT_TAG_NAME = "MARIMO-UI-ELEMENT";
+
+interface IUIElement extends HTMLElement {
+  reset(): void;
+}
+
+/**
+ * Lazily initialize the UIElement component.
+ */
+export function initializeUIElement() {
+  /**
+   * UIElement Web Component
+   *
+   * Synchronizes the value of its first child on page and with the kernel.
+   *
+   * @example
+   * Example wrapping a custom web component:
+   * ```
+   * <marimo-ui-element object-id="...">
+   *   <my-custom-component data-initial-value="..."/>
+   * </marimo-ui-element>
+   * ```
+   *
+   * @example
+   * Example wrapping raw HTML:
+   * ```
+   * <marimo-ui-element object-id="...">
+   *   <div data-initial-value="...">
+   *      ...
+   *   </div>
+   * </marimo-ui-element>
+   * ```
+   *
+   * @remarks
+   * IDENTIFICATION
+   * UIElements are uniquely identified by their objectId, in the following
+   * sense: UIElements with the same objectId are synchronized to have the same
+   * value. In other words, the set of UIElements on the page can be partitioned
+   * into equivalence classes, where two elements are equivalent if they share
+   * the same objectId.
+   *
+   * Every UIElement is registered with the global UIElementRegistry.
+   *
+   * SYNCHRONIZATION
+   * Using a <marimo-ui-element> tag declares that its first child is a component
+   * whose value should be synchronized. Synchronization happens on two levels:
+   *   1. multiple instances of the same component are synchronized to have the
+   *      same value
+   *   2. a change in value on the page is sent to the kernel
+   *
+   * INITIAL VALUE
+   * The first child of <marimo-ui-element> may optionally take a data attribute
+   * called "initial-value" to influence its instantiation. Upon instantiation,
+   * the UIElement component may set this attribute if it already has a value
+   * for the objectId.
+   *
+   * COMMUNICATION
+   * Communication between marimo and the child of a UIElement is facilitated
+   * by two events:
+   *   1. MarimoValueInputEvent: the child publishes its value by dispatching
+   *      a MarimoValueInputEvent, with `detail` set to
+   *      `{ value: <the new value> }`;
+   *   2. MarimoValueUpdateEvent: the UIElement node broadcasts a
+   *      MarimoValueUpdateEvent to every element registered under its
+   *      objectId when a new value is available; the targets of these
+   *      events (i.e., the child node of each UIElement in the equivalence
+   *      class of the objectId) are responsible for listening to these
+   *      events and updating their value internally.
+   *
+   */
+  class UIElement extends HTMLElement implements IUIElement {
+    private initialized = false;
+    private inputListener: (e: MarimoValueInputEventType) => void =
+      Functions.NOOP;
+
+    // This needs to happen in connectedCallback because the element may not be
+    // set at construction time
+    private init() {
+      if (this.initialized) {
+        return;
+      }
+
+      const objectId = UIElementId.parseOrThrow(this);
+      this.inputListener = (e: MarimoValueInputEventType) => {
+        // TODO: just fill in the objectId and let the document handle
+        // broadcast? that would still let other elements cancel the event
+        // while also reducing the number of event listeners on the document
+        if (objectId !== null && e.detail.element === this.firstElementChild) {
+          // A UIElement may be missing from the registry if it was returned from a function that caches return values.
+          if (!UI_ELEMENT_REGISTRY.has(objectId)) {
+            UI_ELEMENT_REGISTRY.registerInstance(
+              objectId,
+              child as HTMLElement,
+            );
+          }
+
+          UI_ELEMENT_REGISTRY.broadcastValueUpdate(
+            child as HTMLElement,
+            objectId,
+            e.detail.value,
+          );
+        }
+      };
+
+      // A UIElement tracks the value of its first child.
+      const child = this.firstElementChild;
+      if (objectId === null) {
+        Logger.error("[marimo-ui-element] missing object-id attribute");
+        return;
+      }
+      if (child === null) {
+        Logger.error("[marimo-ui-element] has no child");
+        return;
+      }
+      if (!(child instanceof HTMLElement)) {
+        Logger.error(
+          "[marimo-ui-element] first child must be instance of HTMLElement",
+        );
+        return;
+      }
+
+      this.initialized = true;
+    }
+
+    connectedCallback() {
+      this.init();
+
+      if (this.initialized) {
+        // It is critical that the element is registered in this method,
+        // and not in the constructor, since it may be disconnected and
+        // reconnected without being re-constructed
+        const objectId = UIElementId.parseOrThrow(this);
+        const child = this.firstElementChild as HTMLElement;
+        UI_ELEMENT_REGISTRY.registerInstance(objectId, child);
+
+        // Listen to marimo input events provided by the child element: these
+        // events are signals to this UIElement that our value should change.
+        document.addEventListener(
+          MarimoValueInputEvent.TYPE,
+          this.inputListener,
+        );
+      }
+    }
+
+    disconnectedCallback() {
+      if (this.initialized) {
+        // Unregister everything
+        document.removeEventListener(
+          MarimoValueInputEvent.TYPE,
+          this.inputListener,
+        );
+        const objectId = UIElementId.parseOrThrow(this);
+        UI_ELEMENT_REGISTRY.removeInstance(
+          objectId,
+          this.firstElementChild as HTMLElement,
+        );
+      }
+    }
+
+    /**
+     * Reset the value of the child element to its initial value.
+     */
+    reset() {
+      const child = this.firstElementChild;
+      if (isCustomMarimoElement(child)) {
+        child.reset();
+      } else {
+        Logger.error(
+          "[marimo-ui-element] first child must have a reset method",
+        );
+      }
+    }
+
+    // We look for changes to the random-id attribute, which is effectively
+    // used like a React key. If the random-id changes, we need to unmount and
+    // remount its child.
+    static get observedAttributes() {
+      return ["random-id"];
+    }
+
+    attributeChangedCallback(
+      name: string,
+      oldValue: string | null,
+      newValue: string | null,
+    ) {
+      if (this.initialized) {
+        const hasChanged = oldValue !== newValue;
+        if (name === "random-id" && hasChanged) {
+          // deregister/clean-up this instance
+          this.disconnectedCallback();
+          // remove and re-add its child to force it to re-render; note that
+          // this doesn't reconstruct the UI element, only its child
+          const child = this.firstElementChild;
+          if (isCustomMarimoElement(child)) {
+            child.rerender();
+          } else {
+            Logger.error(
+              "[marimo-ui-element] first child must have a rerender method",
+            );
+          }
+          // register the element and reset its initial value
+          this.initialized = false;
+          this.connectedCallback();
+        }
+      }
+    }
+  }
+
+  defineCustomElement(UI_ELEMENT_TAG_NAME.toLowerCase(), UIElement);
+}
+
+/**
+ * Given a node, check if its parent or itself is a UIElement,
+ * and return its objectId if so.
+ */
+export function getUIElementObjectId(target: HTMLElement): UIElementId | null {
+  if (!target) {
+    return null;
+  }
+
+  if (target.nodeName === UI_ELEMENT_TAG_NAME) {
+    return UIElementId.parseOrThrow(target);
+  }
+
+  const node = target.parentElement;
+  if (node?.nodeName === UI_ELEMENT_TAG_NAME) {
+    return UIElementId.parseOrThrow(node);
+  }
+
+  return null;
+}
+
+export function isUIElement(target: HTMLElement): target is IUIElement {
+  return target.tagName === UI_ELEMENT_TAG_NAME;
+}

package/src/core/dom/uiregistry.ts

@@ -0,0 +1,209 @@
+/* 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 {
+  MarimoIncomingMessageEvent,
+  MarimoValueReadyEvent,
+  MarimoValueUpdateEvent,
+  type ValueType,
+} from "./events";
+import { parseInitialValue } from "./htmlUtils";
+
+interface UIElementEntry {
+  objectId: string;
+  value: ValueType;
+  // elements synchronized by a UIElement (in the DOM, each
+  // <marimo-ui-element>'s first child)
+  elements: Set<HTMLElement>;
+}
+
+/* UIElementRegistry
+ *
+ * Maintains state for each UI element that is rendered in the
+ * DOM (including shadow DOMs), and handles dispatching of events
+ * to UI elements when values are updated.
+ *
+ * Expects object IDs to have the form <cellId>-<suffixId>
+ */
+export class UIElementRegistry {
+  // maps UIElement objectIds to entries.
+  entries: Map<string, UIElementEntry>;
+
+  /**
+   * Shared instance of UIElementRegistry since this must be a singleton.
+   */
+  static get INSTANCE(): UIElementRegistry {
+    const KEY = "_marimo_private_UIElementRegistry";
+    if (!window[KEY]) {
+      window[KEY] = new UIElementRegistry();
+    }
+    return window[KEY] as UIElementRegistry;
+  }
+
+  private constructor() {
+    this.entries = new Map();
+  }
+
+  has(objectId: UIElementId): boolean {
+    return this.entries.has(objectId);
+  }
+
+  set(objectId: UIElementId, value: ValueType): void {
+    if (this.entries.has(objectId)) {
+      throw new Error(`UIElement ${objectId} already registered`);
+    }
+    this.entries.set(objectId, {
+      objectId: objectId,
+      value: value,
+      elements: new Set(),
+    });
+  }
+
+  /**
+   * Register an instance of a UIElement
+   *
+   * @param objectId - id of the UIElement
+   * @param instance - the HTMLElement that the UIElement wraps
+   */
+  registerInstance(objectId: UIElementId, instance: HTMLElement) {
+    const entry = this.entries.get(objectId);
+    if (entry === undefined) {
+      this.entries.set(objectId, {
+        objectId: objectId,
+        value: parseInitialValue(instance, this),
+        elements: new Set([instance]),
+      });
+    } else {
+      entry.elements.add(instance);
+    }
+  }
+
+  /**
+   * Remove an instance of a UIElement
+   *
+   * @remarks
+   * Should be called when a UIElement node is removed from the DOM.
+   *
+   * @param objectId - id of the UIElement
+   * @param instance - the HTMLElement to remove
+   *
+   */
+  removeInstance(objectId: UIElementId, instance: HTMLElement) {
+    const entry = this.entries.get(objectId);
+    // The UIElement can be removed from the registry before all
+    // instances are removed: UIElement removal is triggered
+    // when the tied Python object goes out of scope, but instance
+    // removal is triggered when the instance is removed from the DOM
+    // (which can happen after Python deconstruction)
+    if (entry?.elements.has(instance)) {
+      entry.elements.delete(instance);
+    }
+  }
+
+  /**
+   * Remove all UIElements associated with a particular cell from the registry.
+   *
+   * Doesn't destroy or unmount HTML elements, just removes associated state
+   * from the registry.
+   *
+   * @param cellId - stringified cellId
+   */
+  removeElementsByCell(cellId: CellId) {
+    const objectIds = [...this.entries.keys()].filter((objectId) =>
+      objectId.startsWith(`${cellId}-`),
+    );
+
+    objectIds.forEach((objectId) => {
+      this.entries.delete(objectId);
+    });
+  }
+
+  /**
+   * Get the value of a registered UIElement.
+   *
+   * @param objectId - id of the UIElement
+   * @returns the value for `objectId`, or `undefined` if the object was not found.
+   */
+  lookupValue(objectId: string): ValueType {
+    const entry = this.entries.get(objectId);
+    return entry === undefined ? undefined : entry.value;
+  }
+
+  broadcastMessage(
+    objectId: UIElementId,
+    message: unknown,
+    buffers: Base64String[] | undefined | null,
+  ): 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({
+            bubbles: false, // only the intended target gets the message
+            composed: true,
+            detail: {
+              objectId: objectId,
+              message: message,
+              buffers: buffers ? buffers.map(toDataView) : undefined,
+            },
+          }),
+        );
+      });
+    }
+  }
+
+  /**
+   * Broadcast `value` to instances of the component with id `objectId`
+   *
+   * Additionally, sends a message to alert the app that an object has a new
+   * value that should be sent to the kernel.
+   *
+   * @param initiator - child element that initiated the broadcast
+   * @param objectId - id of the UIElement
+   * @param value - value to broadcast
+   */
+  broadcastValueUpdate(
+    initiator: HTMLElement,
+    objectId: UIElementId,
+    value: ValueType,
+  ): void {
+    const entry = this.entries.get(objectId);
+    if (entry === undefined) {
+      Logger.warn("UIElementRegistry missing entry", objectId);
+    } else {
+      entry.value = value;
+      entry.elements.forEach((element) => {
+        if (element !== initiator) {
+          element.dispatchEvent(
+            MarimoValueUpdateEvent.create({
+              bubbles: false, // only the intended target gets the message
+              composed: true,
+              detail: { value: value, element: element },
+            }),
+          );
+        }
+      });
+
+      document.dispatchEvent(
+        MarimoValueReadyEvent.create({
+          bubbles: true,
+          composed: true,
+          detail: {
+            objectId: objectId,
+          },
+        }),
+      );
+    }
+  }
+}
+
+export const UI_ELEMENT_REGISTRY = UIElementRegistry.INSTANCE;

package/src/core/edit-app.tsx

@@ -0,0 +1,189 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import { usePrevious } from "@dnd-kit/utilities";
+import { TooltipProvider } from "@radix-ui/react-tooltip";
+import { useAtomValue, useSetAtom } from "jotai";
+import { useEffect } from "react";
+import { Controls } from "@/components/editor/controls/Controls";
+import { AppHeader } from "@/components/editor/header/app-header";
+import { FilenameForm } from "@/components/editor/header/filename-form";
+import { MultiCellActionToolbar } from "@/components/editor/navigation/multi-cell-action-toolbar";
+import { cn } from "@/utils/cn";
+import { Paths } from "@/utils/paths";
+import { AppContainer } from "../components/editor/app-container";
+import {
+  useRunAllCells,
+  useRunStaleCells,
+} from "../components/editor/cell/useRunCells";
+import { CellArray } from "../components/editor/renderers/CellArray";
+import { CellsRenderer } from "../components/editor/renderers/cells-renderer";
+import { useHotkey } from "../hooks/useHotkey";
+import {
+  cellIdsAtom,
+  hasCellsAtom,
+  notebookIsRunningAtom,
+  numColumnsAtom,
+  useCellActions,
+} from "./cells/cells";
+import { CellEffects } from "./cells/effects";
+import type { AppConfig, UserConfig } from "./config/config-schema";
+import { RuntimeState } from "./kernel/RuntimeState";
+import { getSessionId } from "./kernel/session";
+import { useTogglePresenting } from "./layout/useTogglePresenting";
+import { viewStateAtom } from "./mode";
+import { useRequestClient } from "./network/requests";
+import { useFilename } from "./saving/filename";
+import { lastSavedNotebookAtom } from "./saving/state";
+import { useJotaiEffect } from "./state/jotai";
+import { useMarimoWebSocket } from "./websocket/useMarimoWebSocket";
+
+interface AppProps {
+  /**
+   * The user config.
+   */
+  userConfig: UserConfig;
+  /**
+   * The app config.
+   */
+  appConfig: AppConfig;
+  /**
+   * If true, the floating controls will be hidden.
+   */
+  hideControls?: boolean;
+}
+
+export const EditApp: React.FC<AppProps> = ({
+  userConfig,
+  appConfig,
+  hideControls = false,
+}) => {
+  useJotaiEffect(cellIdsAtom, CellEffects.onCellIdsChange);
+
+  const { setCells, mergeAllColumns, collapseAllCells, expandAllCells } =
+    useCellActions();
+  const viewState = useAtomValue(viewStateAtom);
+  const numColumns = useAtomValue(numColumnsAtom);
+  const hasCells = useAtomValue(hasCellsAtom);
+  const filename = useFilename();
+  const setLastSavedNotebook = useSetAtom(lastSavedNotebookAtom);
+  const { sendComponentValues, sendInterrupt } = useRequestClient();
+
+  const isEditing = viewState.mode === "edit";
+  const isPresenting = viewState.mode === "present";
+  const isRunning = useAtomValue(notebookIsRunningAtom);
+
+  // Initialize RuntimeState event-listeners
+  useEffect(() => {
+    RuntimeState.INSTANCE.start(sendComponentValues);
+    return () => {
+      RuntimeState.INSTANCE.stop();
+    };
+  }, []);
+
+  const { connection } = useMarimoWebSocket({
+    autoInstantiate: userConfig.runtime.auto_instantiate,
+    setCells: (cells, layout) => {
+      setCells(cells);
+      const names = cells.map((cell) => cell.name);
+      const codes = cells.map((cell) => cell.code);
+      const configs = cells.map((cell) => cell.config);
+      setLastSavedNotebook({ names, codes, configs, layout });
+    },
+    sessionId: getSessionId(),
+  });
+
+  // Update document title whenever filename or app_title changes
+  useEffect(() => {
+    // Set document title: app_title takes precedence, then filename, then default
+    document.title =
+      appConfig.app_title ||
+      Paths.basename(filename ?? "") ||
+      "Untitled Notebook";
+  }, [appConfig.app_title, filename]);
+
+  // Delete column breakpoints if app width changes from "columns"
+  const previousWidth = usePrevious(appConfig.width);
+  useEffect(() => {
+    if (previousWidth === "columns" && appConfig.width !== "columns") {
+      mergeAllColumns();
+    }
+  }, [appConfig.width, previousWidth, mergeAllColumns, numColumns]);
+
+  const runStaleCells = useRunStaleCells();
+  const runAllCells = useRunAllCells();
+  const togglePresenting = useTogglePresenting();
+
+  // HOTKEYS
+  useHotkey("global.runStale", () => {
+    runStaleCells();
+  });
+  useHotkey("global.interrupt", () => {
+    sendInterrupt();
+  });
+  useHotkey("global.hideCode", () => {
+    togglePresenting();
+  });
+  useHotkey("global.runAll", () => {
+    runAllCells();
+  });
+  useHotkey("global.collapseAllSections", () => {
+    collapseAllCells();
+  });
+  useHotkey("global.expandAllSections", () => {
+    expandAllCells();
+  });
+
+  const editableCellsArray = (
+    <CellArray
+      mode={viewState.mode}
+      userConfig={userConfig}
+      appConfig={appConfig}
+    />
+  );
+
+  return (
+    <>
+      <AppContainer
+        connection={connection}
+        isRunning={isRunning}
+        width={appConfig.width}
+      >
+        <AppHeader
+          connection={connection}
+          className={cn(
+            "pt-4 sm:pt-12 pb-2 mb-4 print:hidden z-50",
+            // Keep the header sticky when scrolling horizontally, for column mode
+            "sticky left-0",
+          )}
+        >
+          {isEditing && (
+            <div className="flex items-center justify-center container">
+              <FilenameForm filename={filename} />
+            </div>
+          )}
+        </AppHeader>
+
+        {/* Don't render until we have a single cell */}
+        {hasCells && (
+          <CellsRenderer appConfig={appConfig} mode={viewState.mode}>
+            {editableCellsArray}
+          </CellsRenderer>
+        )}
+      </AppContainer>
+      <MultiCellActionToolbar />
+      {!hideControls && (
+        <TooltipProvider>
+          <Controls
+            presenting={isPresenting}
+            onTogglePresenting={togglePresenting}
+            onInterrupt={sendInterrupt}
+            onRun={runStaleCells}
+            connectionState={connection.state}
+            running={isRunning}
+            appConfig={appConfig}
+          />
+        </TooltipProvider>
+      )}
+    </>
+  );
+};

package/src/core/errors/__tests__/errors.test.ts

@@ -0,0 +1,62 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+import { describe, expect, it } from "vitest";
+import type { MarimoError } from "@/core/kernel/messages";
+import { getAutoFixes, getImportCode } from "../errors";
+
+describe("getImportCode", () => {
+  it("returns simple import for same name", () => {
+    expect(getImportCode("json")).toBe("import json");
+    expect(getImportCode("math")).toBe("import math");
+  });
+
+  it("returns aliased import for different names", () => {
+    expect(getImportCode("np")).toBe("import numpy as np");
+    expect(getImportCode("pd")).toBe("import pandas as pd");
+    expect(getImportCode("plt")).toBe("import matplotlib.pyplot as plt");
+  });
+});
+
+describe("getAutoFixes", () => {
+  it("returns wrap in function fix for multiple-defs error", () => {
+    const error: MarimoError = {
+      type: "multiple-defs",
+      name: "foo",
+      cells: ["foo"],
+    };
+
+    const fixes = getAutoFixes(error);
+    expect(fixes).toHaveLength(1);
+    expect(fixes[0].title).toBe("Fix: Wrap in a function");
+  });
+
+  it("returns import fix for NameError with known import", () => {
+    const error: MarimoError = {
+      type: "exception",
+      exception_type: "NameError",
+      msg: "name 'np' is not defined",
+    };
+
+    const fixes = getAutoFixes(error);
+    expect(fixes).toHaveLength(1);
+    expect(fixes[0].title).toBe("Fix: Add 'import numpy as np'");
+  });
+
+  it("returns no fixes for NameError with unknown import", () => {
+    const error: MarimoError = {
+      type: "exception",
+      exception_type: "NameError",
+      msg: "name 'unknown_module' is not defined",
+    };
+
+    expect(getAutoFixes(error)).toHaveLength(0);
+  });
+
+  it("returns no fixes for other error types", () => {
+    const error: MarimoError = {
+      type: "syntax",
+      msg: "invalid syntax",
+    };
+
+    expect(getAutoFixes(error)).toHaveLength(0);
+  });
+});