feat: include local wheels in WASM exports

frontend/src/__tests__/main.test.tsx

@@ -16,6 +16,7 @@ import { marimoVersionAtom, showCodeInRunModeAtom } from "../core/meta/state";
 import { type AppMode, viewStateAtom } from "../core/mode";
 import { codeAtom, filenameAtom } from "../core/saving/file-state";
 import { store } from "../core/state/jotai";
+import { wasmWheelUrlsAtom } from "../core/wasm/state";
 import { mount, visibleForTesting } from "../mount";
 
 vi.mock("../utils/vitals", () => ({
@@ -43,6 +44,7 @@ describe("main", () => {
     store.set(appConfigAtom, parseAppConfig({}));
     store.set(userConfigAtom, defaultUserConfig());
     store.set(configOverridesAtom, {});
+    store.set(wasmWheelUrlsAtom, []);
   });
 
   it.each(["edit", "read", "home", "run"])(
@@ -133,6 +135,7 @@ describe("main", () => {
       configOverrides: { display: { code_editor_font_size: 100 } },
       appConfig: { app_title: "My App" } as AppConfig,
       view: { showAppCode: true },
+      wasmWheelUrls: ["public/wheels/demo_pkg-0.1.0-py3-none-any.whl"],
     };
 
     mount(options, el);
@@ -155,6 +158,9 @@ describe("main", () => {
     expect(store.get(appConfigAtom)).toEqual(
       expect.objectContaining({ app_title: "My App" }),
     );
+    expect(store.get(wasmWheelUrlsAtom)).toEqual([
+      "public/wheels/demo_pkg-0.1.0-py3-none-any.whl",
+    ]);
   });
 
   it("should throw on invalid options", () => {

frontend/src/core/wasm/__tests__/bridge.test.ts

@@ -2,8 +2,15 @@
 
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 
-const { mockNotebookReadFile, rpcListeners } = vi.hoisted(() => ({
+const {
+  mockFallbackReadFile,
+  mockNotebookReadFile,
+  mockStartSession,
+  rpcListeners,
+} = vi.hoisted(() => ({
+  mockFallbackReadFile: vi.fn(),
   mockNotebookReadFile: vi.fn(),
+  mockStartSession: vi.fn(),
   rpcListeners: {} as Record<string, () => void>,
 }));
 
@@ -29,7 +36,7 @@ vi.mock("@/core/wasm/rpc", () => ({
     proxy: {
       request: {
         bridge: vi.fn(),
-        startSession: vi.fn(),
+        startSession: mockStartSession,
         readFile: vi.fn(),
         readNotebook: vi.fn(),
         saveNotebook: vi.fn(),
@@ -51,14 +58,18 @@ vi.mock("@/core/wasm/utils", () => ({
 }));
 
 vi.mock("@/core/wasm/store", () => ({
-  fallbackFileStore: { readFile: vi.fn(), saveFile: vi.fn() },
+  fallbackFileStore: { readFile: mockFallbackReadFile, saveFile: vi.fn() },
   notebookFileStore: { readFile: mockNotebookReadFile, saveFile: vi.fn() },
 }));
 
 // Import after all mocks are set up
+import { defaultUserConfig } from "@/core/config/config-schema";
+import { userConfigAtom } from "@/core/config/config";
 import { store } from "@/core/state/jotai";
 import { initialModeAtom } from "@/core/mode";
+import { filenameAtom } from "@/core/saving/file-state";
 import { getWasmWorkerName, PyodideBridge } from "../bridge";
+import { wasmWheelUrlsAtom } from "../state";
 
 // Access INSTANCE once at module level so the constructor runs (and
 // addMessageListener populates rpcListeners) before any test executes.
@@ -67,10 +78,45 @@ void PyodideBridge.INSTANCE;
 describe("PyodideBridge.readCode", () => {
   beforeEach(() => {
     vi.clearAllMocks();
+    mockFallbackReadFile.mockResolvedValue("");
+    mockNotebookReadFile.mockResolvedValue("");
+    mockStartSession.mockResolvedValue(undefined);
+    store.set(filenameAtom, "notebook.py");
+    store.set(initialModeAtom, "read");
+    store.set(userConfigAtom, defaultUserConfig());
+    store.set(wasmWheelUrlsAtom, []);
   });
 
   afterEach(() => {
+    store.set(filenameAtom, null);
     store.set(initialModeAtom, undefined);
+    store.set(userConfigAtom, defaultUserConfig());
+    store.set(wasmWheelUrlsAtom, []);
+  });
+
+  it("passes included wheel URLs to the worker", async () => {
+    mockNotebookReadFile.mockResolvedValue("import demo_pkg");
+    store.set(wasmWheelUrlsAtom, [
+      "public/wheels/demo_pkg-0.1.0-py3-none-any.whl",
+      "https://cdn.example.com/extra_pkg-0.1.0-py3-none-any.whl",
+    ]);
+
+    rpcListeners.ready();
+    await new Promise((resolve) => setTimeout(resolve, 0));
+
+    expect(mockStartSession).toHaveBeenCalledWith(
+      expect.objectContaining({
+        code: "import demo_pkg",
+        filename: "notebook.py",
+        wheelUrls: [
+          new URL(
+            "public/wheels/demo_pkg-0.1.0-py3-none-any.whl",
+            document.baseURI,
+          ).toString(),
+          "https://cdn.example.com/extra_pkg-0.1.0-py3-none-any.whl",
+        ],
+      }),
+    );
   });
 
   it("reads from notebookFileStore in read mode", async () => {

frontend/src/core/wasm/bridge.ts

@@ -38,7 +38,11 @@ import type { IConnectionTransport } from "../websocket/transports/transport";
 import { PyodideRouter } from "./router";
 import { getWorkerRPC } from "./rpc";
 import { createShareableLink } from "./share";
-import { wasmInitializationAtom, wasmInitStatusAtom } from "./state";
+import {
+  wasmInitializationAtom,
+  wasmInitStatusAtom,
+  wasmWheelUrlsAtom,
+} from "./state";
 import { fallbackFileStore, notebookFileStore } from "./store";
 import { isWasm } from "./utils";
 import type { SaveWorkerSchema } from "./worker/save-worker";
@@ -155,6 +159,9 @@ export class PyodideBridge implements RunRequests, EditRequests {
     const fallbackCode = await fallbackFileStore.readFile();
     const filename = store.get(filenameAtom) ?? PyodideRouter.getFilename();
     const userConfig = store.get(userConfigAtom);
+    const wheelUrls = store
+      .get(wasmWheelUrlsAtom)
+      .map((url) => new URL(url, document.baseURI).toString());
 
     const queryParameters: Record<string, string | string[]> = {};
     const searchParams = new URLSearchParams(window.location.search);
@@ -167,6 +174,7 @@ export class PyodideBridge implements RunRequests, EditRequests {
       queryParameters: queryParameters,
       code: code || fallbackCode || "",
       filename,
+      wheelUrls,
       userConfig: {
         ...userConfig,
         runtime: {

frontend/src/core/wasm/state.ts

@@ -8,6 +8,8 @@ export const wasmInitializationAtom = atom<string>("Initializing...");
 export type WasmInitStatus = "loading" | "ready" | "error";
 export const wasmInitStatusAtom = atom<WasmInitStatus>("loading");
 
+export const wasmWheelUrlsAtom = atom<string[]>([]);
+
 export const hasAnyOutputAtom = atom<boolean>((get) => {
   const notebook = get(notebookAtom);
   const runtimeStates = Object.values(notebook.cellRuntime);

frontend/src/core/wasm/worker/bootstrap.ts

@@ -108,10 +108,18 @@ export class DefaultWasmController implements WasmController {
     queryParameters: Record<string, string | string[]>;
     code: string;
     filename: string | null;
+    wheelUrls?: string[];
     onMessage: (message: JsonString<NotificationPayload>) => void;
     userConfig: UserConfig;
   }): Promise<SerializedBridge> {
-    const { code, filename, onMessage, queryParameters, userConfig } = opts;
+    const {
+      code,
+      filename,
+      onMessage,
+      queryParameters,
+      userConfig,
+      wheelUrls = [],
+    } = opts;
     // We pass down a messenger object to the code
     // This is used to have synchronous communication between the JS and Python code
     // Previously, we used a queue, but this would not properly flush the queue
@@ -124,6 +132,8 @@ export class DefaultWasmController implements WasmController {
     self.query_params = queryParameters;
     self.user_config = userConfig;
 
+    await this.installIncludedWheels(wheelUrls);
+
     const span = t.startSpan("startSession.runPython");
     const nbFilename = filename || WasmFileSystem.NOTEBOOK_FILENAME;
     const [bridge, init, packages] = this.requirePyodide.runPython(
@@ -169,6 +179,23 @@ export class DefaultWasmController implements WasmController {
     return bridge;
   }
 
+  private async installIncludedWheels(wheelUrls: string[]) {
+    if (wheelUrls.length === 0) {
+      return;
+    }
+
+    const loadSpan = t.startSpan("micropip.install.wheels");
+    try {
+      await this.requirePyodide.runPythonAsync(`
+        import micropip
+        print("Loading included wheels:", ${JSON.stringify(wheelUrls)})
+        await micropip.install(${JSON.stringify(wheelUrls)})
+      `);
+    } finally {
+      loadSpan.end();
+    }
+  }
+
   private async loadNotebookDeps(code: string, foundPackages: Set<string>) {
     const pyodide = this.requirePyodide;
 

frontend/src/core/wasm/worker/types.ts

@@ -62,6 +62,7 @@ export interface WasmController {
     queryParameters: Record<string, string | string[]>;
     code: string;
     filename: string | null;
+    wheelUrls?: string[];
     userConfig: UserConfig;
     onMessage: (message: JsonString<NotificationPayload>) => void;
   }): Promise<SerializedBridge>;

frontend/src/core/wasm/worker/worker.ts

@@ -86,6 +86,7 @@ const requestHandler = createRPCRequestHandler({
     queryParameters: Record<string, string | string[]>;
     code: string;
     filename: string | null;
+    wheelUrls?: string[];
     userConfig: UserConfig;
   }) => {
     await pyodideReadyPromise; // Make sure loading is done
@@ -118,6 +119,7 @@ const requestHandler = createRPCRequestHandler({
         code: opts.code,
         filename: opts.filename,
         queryParameters: opts.queryParameters,
+        wheelUrls: opts.wheelUrls,
         userConfig: opts.userConfig,
         onMessage: (msg) => {
           initializeOnce();

frontend/src/mount.tsx

@@ -51,6 +51,7 @@ import {
 import { maybeRegisterVSCodeBindings } from "./core/vscode/vscode-bindings";
 import type { FileStore } from "./core/wasm/store";
 import { notebookFileStore } from "./core/wasm/store";
+import { wasmWheelUrlsAtom } from "./core/wasm/state";
 import { WebSocketState } from "./core/websocket/types";
 import {
   handleWidgetMessage,
@@ -269,6 +270,14 @@ const mountOptionsSchema = z.object({
     )
     .nullish()
     .transform((val) => val ?? []),
+
+  /**
+   * Local wheel URLs bundled with a WASM export.
+   */
+  wasmWheelUrls: z
+    .array(z.string())
+    .nullish()
+    .transform((val) => val ?? []),
 });
 
 function initStore(options: unknown) {
@@ -309,6 +318,7 @@ function initStore(options: unknown) {
   // Meta
   store.set(marimoVersionAtom, parsedOptions.data.version);
   store.set(showCodeInRunModeAtom, parsedOptions.data.view.showAppCode);
+  store.set(wasmWheelUrlsAtom, parsedOptions.data.wasmWheelUrls);
 
   // Check for view-as parameter to start in present mode
   const shouldStartInPresentMode = (() => {

marimo/_cli/export/commands.py

@@ -7,6 +7,7 @@
 from dataclasses import replace
 from pathlib import Path
 from typing import TYPE_CHECKING, Literal
+from urllib.parse import quote
 
 import click
 
@@ -66,6 +67,8 @@
     "`uv run --isolated`. Requires `uv`."
 )
 
+_WASM_WHEEL_URL_PREFIX = "public/wheels"
+
 
 @click.group(
     cls=ColoredGroup, help="""Export a notebook to various formats."""
@@ -74,6 +77,40 @@ def export() -> None:
     pass
 
 
+def _included_wheel_urls(wheel_paths: tuple[Path, ...]) -> list[str]:
+    wheel_urls: list[str] = []
+    seen_names: set[str] = set()
+    for wheel_path in wheel_paths:
+        if wheel_path.suffix.lower() != ".whl":
+            raise click.UsageError(
+                f"--include-wheel expects a .whl file: {wheel_path}"
+            )
+        if wheel_path.name in seen_names:
+            raise click.UsageError(
+                "--include-wheel files must have distinct filenames: "
+                f"{wheel_path.name}"
+            )
+        seen_names.add(wheel_path.name)
+        wheel_urls.append(f"{_WASM_WHEEL_URL_PREFIX}/{quote(wheel_path.name)}")
+    return wheel_urls
+
+
+def _copy_included_wheels(
+    out_dir: Path, wheel_paths: tuple[Path, ...]
+) -> None:
+    if not wheel_paths:
+        return
+
+    import shutil
+
+    wheel_dir = out_dir / _WASM_WHEEL_URL_PREFIX
+    wheel_dir.mkdir(parents=True, exist_ok=True)
+    for wheel_path in wheel_paths:
+        target = wheel_dir / wheel_path.name
+        if wheel_path.resolve() != target.resolve():
+            shutil.copyfile(wheel_path, target)
+
+
 class _PDFExportCLIReporter:
     def __call__(self, event: PDFExportStatusEvent) -> None:
         message = event.message
@@ -890,6 +927,21 @@ def export_callback_impl(file_path: MarimoPath) -> ExportResult:
         " (index.js and wrangler.jsonc) for easy deployment."
     ),
 )
+@click.option(
+    "--include-wheel",
+    "include_wheels",
+    type=click.Path(
+        exists=True,
+        file_okay=True,
+        dir_okay=False,
+        path_type=Path,
+    ),
+    multiple=True,
+    help=(
+        "Include a local Python wheel in the WASM export. "
+        "Can be passed multiple times."
+    ),
+)
 @click.option(
     "--sandbox/--no-sandbox",
     is_flag=True,
@@ -926,6 +978,7 @@ def html_wasm(
     watch: bool,
     show_code: bool,
     include_cloudflare: bool,
+    include_wheels: tuple[Path, ...],
     sandbox: bool | None,
     force: bool,
     execute: bool,
@@ -982,6 +1035,7 @@ def html_wasm(
         filename = output.name
 
     marimo_file = MarimoPath(name)
+    wasm_wheel_urls = _included_wheel_urls(include_wheels)
 
     if execute:
         cli_args = parse_args(args)
@@ -1004,14 +1058,20 @@ def export_callback(file_path: MarimoPath) -> ExportResult:
                     show_code=show_code,
                     cli_args=cli_args,
                     argv=list(args),
+                    wasm_wheel_urls=wasm_wheel_urls,
                 )
             )
 
         echo("Executing notebook...")
     else:
 
         def export_callback(file_path: MarimoPath) -> ExportResult:
-            return export_as_wasm(file_path, mode, show_code=show_code)
+            return export_as_wasm(
+                file_path,
+                mode,
+                show_code=show_code,
+                wasm_wheel_urls=wasm_wheel_urls,
+            )
 
     # Export assets first
     Exporter().export_assets(out_dir)
@@ -1031,6 +1091,7 @@ def export_callback(file_path: MarimoPath) -> ExportResult:
             f"The public folder next to your notebook was copied to "
             f"{green(str(out_dir))}."
         )
+    _copy_included_wheels(out_dir, include_wheels)
 
     echo(
         "To run the exported notebook, use:\n"