From 7fb7afe3d7cbde29f00a0d104409f0d4cccdac83 Mon Sep 17 00:00:00 2001
From: Muspi Merol <wjj123345567@icloud.com>
Date: Wed, 16 Oct 2024 15:48:37 +0800
Subject: [PATCH] feat: explore the feasibility of running pyodide inside a
 sharedworker

---
 src/lib/pyodide/start/init.ts       | 10 +++----
 src/lib/pyodide/worker/types.ts     | 14 ++++++++++
 src/lib/pyodide/worker/worker.ts    | 43 +++++++++++++++++++++++++++++
 src/lib/pyodide/worker/workerApi.ts | 33 ++++++++++++++++++++++
 src/python/common/patches.py        | 20 +-------------
 src/python/common/toast.py          |  2 --
 src/routes/test/+page.svelte        | 37 +++++++++++++++++++++++++
 7 files changed, 133 insertions(+), 26 deletions(-)
 create mode 100644 src/lib/pyodide/worker/types.ts
 create mode 100644 src/lib/pyodide/worker/worker.ts
 create mode 100644 src/lib/pyodide/worker/workerApi.ts
 create mode 100644 src/routes/test/+page.svelte

diff --git a/src/lib/pyodide/start/init.ts b/src/lib/pyodide/start/init.ts
index c36c88e0..249785bd 100644
--- a/src/lib/pyodide/start/init.ts
+++ b/src/lib/pyodide/start/init.ts
@@ -6,15 +6,15 @@ import loader from "./loader.py?raw";
 import { dev } from "$app/environment";
 import { cacheSingleton } from "$lib/utils/cache";
 import { getEnv } from "$lib/utils/env";
-import { withToast } from "$lib/utils/toast";
-import { toast } from "svelte-sonner";
+// import { withToast } from "$lib/utils/toast";
+// import { toast } from "svelte-sonner";
 
-const getMinimalPyodide = cacheSingleton(withToast({ loading: "loading pyodide runtime" })(async () => {
+const getMinimalPyodide = cacheSingleton(async () => {
   const { loadPyodide } = await import("pyodide");
   const py = await loadPyodide({ indexURL, env: getEnv(), packages: preloadPackages, args: dev ? [] : ["-O"] });
-  py.globals.set("toast", toast);
+  // py.globals.set("toast", toast);
   return py;
-}));
+});
 
 const getSetupModule = cacheSingleton(async () => {
   const py = await getMinimalPyodide();
diff --git a/src/lib/pyodide/worker/types.ts b/src/lib/pyodide/worker/types.ts
new file mode 100644
index 00000000..c71d31eb
--- /dev/null
+++ b/src/lib/pyodide/worker/types.ts
@@ -0,0 +1,14 @@
+export type UUID = ReturnType<Crypto["randomUUID"]>;
+
+interface BaseTask {
+  id: UUID;
+  type: string;
+  data: any;
+}
+
+export interface EvalTask extends BaseTask {
+  type: "eval";
+  data: string;
+}
+
+export type Task = EvalTask;
diff --git a/src/lib/pyodide/worker/worker.ts b/src/lib/pyodide/worker/worker.ts
new file mode 100644
index 00000000..9b1dd750
--- /dev/null
+++ b/src/lib/pyodide/worker/worker.ts
@@ -0,0 +1,43 @@
+/* eslint-disable no-restricted-globals */
+
+/// <reference lib="webworker" />
+
+import type { Task } from "./types";
+
+import getPy from "..";
+
+(self as unknown as SharedWorkerGlobalScope).addEventListener("connect", (event) => {
+  const [port] = event.ports;
+  port.addEventListener("message", handleMessage);
+  port.start();
+  port.postMessage("hello from worker");
+});
+
+async function handleMessage(event: MessageEvent<Task | object>) {
+  if (typeof event.data === "object" && "id" in event.data) {
+    const port = event.currentTarget as MessagePort;
+    const { id, type, data } = event.data;
+    // dispatch
+    if (type === "eval") {
+      runTask(id, port, handleEval(data));
+    }
+  }
+  else {
+    console.warn(event);
+  }
+}
+
+function runTask(id: string, port: MessagePort, promise: Promise<any>) {
+  promise
+    .then(result => port.postMessage({ id, result }))
+    .catch(error => port.postMessage({ id, error }));
+}
+
+async function handleEval(source: string) {
+  // const { loadPyodide: getPy } = await import("pyodide");
+  // const { getPyodide: getPy } = await import("../start/init");
+
+  const py = await getPy();
+  const result = await py.runPythonAsync(source);
+  return String(result);
+}
diff --git a/src/lib/pyodide/worker/workerApi.ts b/src/lib/pyodide/worker/workerApi.ts
new file mode 100644
index 00000000..12cfb3b6
--- /dev/null
+++ b/src/lib/pyodide/worker/workerApi.ts
@@ -0,0 +1,33 @@
+import type { UUID } from "./types";
+
+import { dev } from "$app/environment";
+import { cacheSingleton } from "$lib/utils/cache";
+
+const getWorker = cacheSingleton(() => {
+  const worker = new SharedWorker(new URL("./worker", import.meta.url), { name: "pyodide", type: /* @vite-ignore */ dev ? "module" : "classic" });
+  worker.port.addEventListener("message", handleMessage);
+  worker.port.start();
+  worker.port.postMessage("hello from main thread");
+  return worker;
+});
+
+const pendingTasks = new Map<UUID, [(value: any) => void, (reason: any) => void]>();
+
+async function handleMessage(event: MessageEvent) {
+  if (typeof event.data === "object" && "id" in event.data) {
+    const data: { id: UUID } & ({ result: any } | { error: any }) = event.data;
+    const [resolve, reject] = pendingTasks.get(data.id)!;
+    "error" in data ? reject(data.error) : resolve(data.result);
+    pendingTasks.delete(data.id);
+  }
+}
+
+export function evalPython(source: string) {
+  const worker = getWorker();
+  const id = crypto.randomUUID();
+  const task = new Promise((resolve, reject) => {
+    pendingTasks.set(id, [resolve, reject]);
+  });
+  worker.port.postMessage({ id, data: source, type: "eval" });
+  return task;
+}
diff --git a/src/python/common/patches.py b/src/python/common/patches.py
index e06c4d2b..d312b3e5 100644
--- a/src/python/common/patches.py
+++ b/src/python/common/patches.py
@@ -1,10 +1,8 @@
-import builtins
 from functools import cache, wraps
-from inspect import Signature, getsource
+from inspect import getsource
 from os import getenv
 
 import micropip
-from js import window
 from pyodide.ffi import can_run_sync, run_sync
 
 from .package import get_package_name
@@ -56,14 +54,6 @@ async def runcode(self: PyodideConsole, source: str, code):
     PyodideConsole.runcode = runcode
 
 
-@cache
-def patch_input():
-    def input(prompt=""):
-        return window.prompt(str(prompt)) or ""
-
-    builtins.input = input
-
-
 @cache
 def patch_sync():
     import asyncio
@@ -95,15 +85,7 @@ def _(duration):
     time.sleep = _
 
 
-@cache
-def patch_exit():
-    window.close.__signature__ = Signature()  # type: ignore
-    builtins.exit = builtins.quit = window.close
-
-
 patch_install()
 patch_linecache()
 patch_console()
-patch_input()
 patch_sync()
-patch_exit()
diff --git a/src/python/common/toast.py b/src/python/common/toast.py
index 85e58aca..be71795d 100644
--- a/src/python/common/toast.py
+++ b/src/python/common/toast.py
@@ -5,8 +5,6 @@
 
 if TYPE_CHECKING:
     from stub import toast
-else:
-    from __main__ import toast
 
 
 class ToastController:
diff --git a/src/routes/test/+page.svelte b/src/routes/test/+page.svelte
new file mode 100644
index 00000000..e233b8fa
--- /dev/null
+++ b/src/routes/test/+page.svelte
@@ -0,0 +1,37 @@
+<script lang="ts">
+  import { evalPython } from "$lib/pyodide/worker/workerApi";
+
+  let input = "";
+  let output: any;
+  let error: any;
+
+  $: if (input) {
+    evalPython(input)
+      .then((res) => {
+        output = res;
+        error = "";
+      })
+      .catch((err) => {
+        error = err;
+        output = "";
+      });
+  }
+  else {
+    output = "";
+    error = "";
+  }
+</script>
+
+<div class="m-10 flex flex-col gap-2.5 font-mono [&>*]:(px-2.5 py-2)">
+
+  <input type="text" bind:value={input} class="bg-neutral-8 outline-none ring-(1.2 neutral-4)">
+
+  {#if error}
+    <div class="ws-pre-wrap text-sm text-red">{error}</div>
+  {/if}
+
+  {#if output}
+    <div class="ws-pre-wrap text-sm text-teal">{output}</div>
+  {/if}
+
+</div>