macroWorker.ts

/**
 * DataLab-Web macro Worker.
 *
 * Loads a second Pyodide instance in a dedicated Web Worker so user
 * macros run isolated from the main UI thread.  Communicates with the
 * main thread via ``postMessage``:
 *
 *   ──► main thread sends:
 *     { type: "init" }
 *     { type: "run", code: string, name?: string }
 *     { type: "bridge_reply", id: string, ok: boolean, value?, error? }
 *
 *   ◄── worker sends:
 *     { type: "ready" }
 *     { type: "stdout"|"stderr", text: string }
 *     { type: "started", name: string }
 *     { type: "finished", ok: boolean, error?: string }
 *     { type: "bridge_call", id: string, method: string, payload: any }
 *
 * Stop is implemented main-side via ``Worker.terminate()``.
 */

import macroProxySource from "./macro_proxy.py?raw";
import dlwTitleFormatSource from "./dlw_title_format.py?raw";
import {
  bootPyodide,
  installBridge,
  installProxyGlobal,
  resolveBridgeReply,
  type DLWWorkerScope,
  type PyodideAPI,
} from "./workerBase";
// Same JSON Schema / backends shims as the main runtime — required so
// ``guidata.dataset`` (transitively imported by ``sigima``) loads
// cleanly under Pyodide. See runtime.ts for the rationale.
import guidataJsonSchemaShim from "./_guidata_jsonschema_shim.py?raw";
const guidataBackendsSource = (() => {
  const candidates = import.meta.glob("./_guidata_backends_shim.py", {
    query: "?raw",
    import: "default",
    eager: true,
  }) as Record<string, string>;
  const first = Object.values(candidates)[0];
  return first ?? null;
})();

declare const self: DedicatedWorkerGlobalScope & DLWWorkerScope;

let pyPromise: Promise<PyodideAPI> | null = null;

// ``LANG`` used to boot Pyodide. Defaults to the POSIX ``C`` (English) locale
// and is overridden by the ``lang`` field of the ``init`` message so the
// macro worker matches the main thread's UI locale (which alone can read
// ``localStorage`` — workers cannot). See ``runtime.ts`` for the rationale.
let pyLang = "C";

/** Load Pyodide on first use; subsequent calls return the same instance. */
async function getPyodide(): Promise<PyodideAPI> {
  if (pyPromise) return pyPromise;
  pyPromise = (async () => {
    const py = await bootPyodide({
      lang: pyLang,
      packages: ["numpy", "scipy", "h5py", "micropip"],
      titleFormatSource: dlwTitleFormatSource,
    });

    await py.runPythonAsync(guidataJsonSchemaShim);
    if (guidataBackendsSource) {
      await py.runPythonAsync(guidataBackendsSource);
    }

    // Stream stdout/stderr to the main thread, line by line.
    // ``setStdout``/``setStderr`` work for raw C-level writes but
    // user-level ``print`` in the auto-wrapped coroutine doesn't always
    // reach them, so we also patch ``sys.stdout`` / ``sys.stderr`` in
    // Python — this is what ultimately captures ``print`` output.
    py.setStdout({
      batched: (s: string) => self.postMessage({ type: "stdout", text: s }),
    });
    py.setStderr({
      batched: (s: string) => self.postMessage({ type: "stderr", text: s }),
    });
    py.globals.set("_dlw_post_stdout", (s: unknown) =>
      self.postMessage({ type: "stdout", text: String(s) }),
    );
    py.globals.set("_dlw_post_stderr", (s: unknown) =>
      self.postMessage({ type: "stderr", text: String(s) }),
    );
    await py.runPythonAsync(`
import sys

class _DLWStream:
    def __init__(self, post):
        self._post = post
        self._buf = ""

    def write(self, s):
        if not s:
            return 0
        self._buf += s
        # Flush complete lines immediately; keep the trailing partial.
        if "\\n" in self._buf:
            head, _, tail = self._buf.rpartition("\\n")
            self._post(head + "\\n")
            self._buf = tail
        return len(s)

    def flush(self):
        if self._buf:
            self._post(self._buf)
            self._buf = ""

    def isatty(self):
        return False

sys.stdout = _DLWStream(_dlw_post_stdout)
sys.stderr = _DLWStream(_dlw_post_stderr)
`);

    // Bridge: macro calls ``js._dlw_bridge_call(method, payload)`` and
    // awaits the returned Promise; we resolve it when the main thread
    // posts back ``{type: "bridge_reply", id, ok, value|error}``.
    installBridge(self, (m) => self.postMessage(m));

    // Install the Python ``proxy`` global by executing macro_proxy.py.
    await installProxyGlobal(py, macroProxySource);
    return py;
  })();
  return pyPromise;
}

// Surface async errors that escape the ``onmessage`` try/catch (e.g. a
// timer scheduled by Pyodide that throws, or a detached Promise that
// rejects).  Without these handlers the worker would die silently — the
// main thread relies on ``stderr`` / ``finished`` messages to update the
// macro console and Stop button state.
self.onerror = (event: Event | string): boolean => {
  const text =
    typeof event === "string"
      ? event
      : (event as ErrorEvent).message || "macro worker error";
  self.postMessage({ type: "stderr", text: text + "\n" });
  self.postMessage({ type: "finished", ok: false, error: text });
  // Returning ``true`` would prevent default logging in the host page;
  // we want the error to also appear in DevTools, so let it propagate.
  return false;
};

self.onunhandledrejection = (event: PromiseRejectionEvent): void => {
  const reason = event.reason;
  const text =
    reason instanceof Error
      ? reason.message
      : typeof reason === "string"
        ? reason
        : "unhandled promise rejection in macro worker";
  self.postMessage({ type: "stderr", text: text + "\n" });
  self.postMessage({ type: "finished", ok: false, error: text });
};

self.onmessage = async (event: MessageEvent) => {
  const msg = event.data as
    | { type: "init"; lang?: string }
    | { type: "run"; code: string; name?: string }
    | {
        type: "bridge_reply";
        id: string;
        ok: boolean;
        value?: unknown;
        error?: string;
      };
  try {
    if (msg.type === "init") {
      if (msg.lang) pyLang = msg.lang;
      await getPyodide();
      self.postMessage({ type: "ready" });
      return;
    }
    if (msg.type === "bridge_reply") {
      resolveBridgeReply(self, msg);
      return;
    }
    if (msg.type === "run") {
      const py = await getPyodide();
      self.postMessage({ type: "started", name: msg.name ?? "" });
      try {
        // ``proxy`` is already a global injected by macro_proxy.py.
        // ``runPythonAsync`` supports top-level ``await`` by auto-
        // wrapping the source in a coroutine.
        await py.runPythonAsync(msg.code);
        // Flush any partial line still buffered in our redirected
        // ``sys.stdout`` / ``sys.stderr`` (no trailing ``\n``).
        await py.runPythonAsync(
          "import sys\nsys.stdout.flush()\nsys.stderr.flush()",
        );
        self.postMessage({ type: "finished", ok: true });
      } catch (err) {
        const text = err instanceof Error ? err.message : String(err);
        self.postMessage({ type: "stderr", text: text + "\n" });
        self.postMessage({ type: "finished", ok: false, error: text });
      }
    }
  } catch (err) {
    const text = err instanceof Error ? err.message : String(err);
    self.postMessage({ type: "stderr", text: text + "\n" });
    self.postMessage({ type: "finished", ok: false, error: text });
  }
};