[next] Enabled hooks around plugins (#1522)

pyscript.core/esm/plugins.js

@@ -1,10 +1,13 @@
 import { $$ } from "basic-devtools";
 
+import { create } from "./utils.js";
 import { getDetails } from "./script-handler.js";
 import { registry, configs } from "./runtimes.js";
 import { getRuntimeID } from "./loader.js";
 import { io } from "./runtime/_utils.js";
 
+import workerHooks from "./worker/hooks.js";
+
 export const PLUGINS_SELECTORS = [];
 
 /**
@@ -18,6 +21,8 @@ export const PLUGINS_SELECTORS = [];
  * @prop {(path:string, data:ArrayBuffer) => void} writeFile an utility to write a file in the virtual FS, if available
  */
 
+const patched = new Map();
+
 // REQUIRES INTEGRATION TEST
 /* c8 ignore start */
 /**
@@ -40,16 +45,80 @@ export const handlePlugin = (node) => {
                     config,
                 );
                 engine.then((runtime) => {
-                    const module = registry.get(type);
-                    onRuntimeReady(node, {
-                        type,
-                        runtime,
-                        XWorker,
-                        io: io.get(runtime),
-                        config: structuredClone(configs.get(name)),
-                        run: module.run.bind(module, runtime),
-                        runAsync: module.runAsync.bind(module, runtime),
-                    });
+                    if (!patched.has(id)) {
+                        const module = create(registry.get(type));
+                        const {
+                            onBeforeRun,
+                            onBeforeRunAsync,
+                            onAfterRun,
+                            onAfterRunAsync,
+                            codeBeforeRunWorker,
+                            codeBeforeRunWorkerAsync,
+                            codeAfterRunWorker,
+                            codeAfterRunWorkerAsync,
+                        } = options;
+
+                        // These two loops mimic a `new Map(arrayContent)` without needing
+                        // the new Map overhead so that [name, [before, after]] can be easily destructured
+                        // and new sync or async patches become easy to add (when the logic is the same).
+
+                        // patch sync
+                        for (const [name, [before, after]] of [
+                            ["run", [onBeforeRun, onAfterRun]],
+                        ]) {
+                            const method = module[name];
+                            module[name] = function (runtime, code) {
+                                if (before) before.call(this, resolved, node);
+                                const result = method.call(this, runtime, code);
+                                if (after) after.call(this, resolved, node);
+                                return result;
+                            };
+                        }
+
+                        // patch async
+                        for (const [name, [before, after]] of [
+                            ["runAsync", [onBeforeRunAsync, onAfterRunAsync]],
+                        ]) {
+                            const method = module[name];
+                            module[name] = async function (runtime, code) {
+                                if (before)
+                                    await before.call(this, resolved, node);
+                                const result = await method.call(
+                                    this,
+                                    runtime,
+                                    code,
+                                );
+                                if (after)
+                                    await after.call(this, resolved, node);
+                                return result;
+                            };
+                        }
+
+                        // setup XWorker hooks, allowing strings to be forwarded to the worker
+                        // whenever it's created, as functions can't possibly be serialized
+                        // unless these are pure with no outer scope access (or globals vars)
+                        // so that making it strings disambiguate about their running context.
+                        workerHooks.set(XWorker, {
+                            beforeRun: codeBeforeRunWorker,
+                            beforeRunAsync: codeBeforeRunWorkerAsync,
+                            afterRun: codeAfterRunWorker,
+                            afterRunAsync: codeAfterRunWorkerAsync,
+                        });
+
+                        const resolved = {
+                            type,
+                            runtime,
+                            XWorker,
+                            io: io.get(runtime),
+                            config: structuredClone(configs.get(name)),
+                            run: module.run.bind(module, runtime),
+                            runAsync: module.runAsync.bind(module, runtime),
+                        };
+
+                        patched.set(id, resolved);
+                    }
+
+                    onRuntimeReady(patched.get(id), node);
                 });
             }
         }

pyscript.core/esm/worker/_template.js

@@ -6,6 +6,7 @@
 
 import coincident from "coincident/structured";
 
+import { create } from "../utils.js";
 import { registry } from "../runtimes.js";
 import { getRuntime, getRuntimeID } from "../loader.js";
 
@@ -55,12 +56,40 @@ const xworker = {
     },
 };
 
-add("message", ({ data: { options, code } }) => {
+add("message", ({ data: { options, code, hooks } }) => {
     engine = (async () => {
         const { type, version, config, async: isAsync } = options;
         const engine = await getRuntime(getRuntimeID(type, version), config);
-        const details = registry.get(type);
-        run = details[`runWorker${isAsync ? "Async" : ""}`].bind(details);
+        const details = create(registry.get(type));
+        const name = `runWorker${isAsync ? "Async" : ""}`;
+
+        // patch code if needed
+        const { beforeRun, beforeRunAsync, afterRun, afterRunAsync } = hooks;
+
+        const after = afterRun || afterRunAsync;
+        const before = beforeRun || beforeRunAsync;
+
+        // append code that should be executed *after* first
+        if (after) {
+            const method = details[name];
+            details[name] = function (runtime, code, xworker) {
+                return method.call(this, runtime, `${code}\n${after}`, xworker);
+            };
+        }
+
+        // prepend code that should be executed *before* (so that after is post-patched)
+        if (before) {
+            const method = details[name];
+            details[name] = function (runtime, code, xworker) {
+                return method.call(
+                    this,
+                    runtime,
+                    `${before}\n${code}`,
+                    xworker,
+                );
+            };
+        }
+        run = details[name].bind(details);
         run(engine, code, xworker);
         return engine;
     })();

pyscript.core/esm/worker/class.js

@@ -2,6 +2,7 @@ import coincident from "coincident/structured";
 import xworker from "./xworker.js";
 import { assign, defineProperties, absoluteURL } from "../utils.js";
 import { getText } from "../fetch-utils.js";
+import workerHooks from "./hooks.js";
 
 /**
  * @typedef {Object} WorkerOptions plugin configuration
@@ -18,6 +19,7 @@ export default (...args) =>
      * @returns {Worker}
      */
     function XWorker(url, options) {
+        const hooks = workerHooks.get(XWorker);
         const worker = xworker();
         const { postMessage } = worker;
         if (args.length) {
@@ -28,7 +30,7 @@ export default (...args) =>
         if (options?.config) options.config = absoluteURL(options.config);
         const bootstrap = fetch(url)
             .then(getText)
-            .then((code) => postMessage.call(worker, { options, code }));
+            .then((code) => postMessage.call(worker, { options, code, hooks }));
         return defineProperties(worker, {
             postMessage: {
                 value: (data, ...rest) =>

pyscript.core/esm/worker/hooks.js

@@ -0,0 +1 @@
+export default new WeakMap();

pyscript.core/package.json

@@ -50,6 +50,6 @@
         "coincident": "^0.2.3"
     },
     "worker": {
-        "blob": "sha256-LhH/RqqHJJNjYkIoMjALOt1xdneBHAyy+jgMKO07Mx0="
+        "blob": "sha256-6voWcbsOeEWD2McAxBkTW89la0C9rF93Fkz5z3+xGGc="
     }
 }

pyscript.core/test/a.py

@@ -1 +1,2 @@
 x = "hello from A"
+print(x)

pyscript.core/test/plugins/index.html

@@ -16,7 +16,7 @@
             import { registerPlugin } from "@pyscript/core";
             registerPlugin("mpy-script", {
                 type: "micropython", // or just 'mpy'
-                async onRuntimeReady(element, micropython) {
+                async onRuntimeReady(micropython, element) {
                     console.log(micropython);
                     // Somehow this doesn't work in MicroPython
                     micropython.io.stdout = (message) => {

pyscript.core/test/plugins/lua.html

@@ -16,7 +16,7 @@
             import { registerPlugin } from "@pyscript/core";
             registerPlugin("lua-script", {
                 type: "wasmoon", // or just 'lua'
-                async onRuntimeReady(element, wasmoon) {
+                async onRuntimeReady(wasmoon, element) {
                     // Somehow this doesn't work in Wasmoon
                     wasmoon.io.stdout = (message) => {
                         console.log("🌑", wasmoon.type, message);

pyscript.core/test/plugins/py-script.html

@@ -28,5 +28,9 @@
             js.console.log(a.x, b.x)
             'Hello Web!'
         </py-script>
+        <py-script>
+            XWorker('../a.py')
+            'OK'
+        </py-script>
     </body>
 </html>

pyscript.core/test/plugins/py-script.js

@@ -11,19 +11,32 @@ document.head.appendChild(document.createElement("style")).textContent = `
 let id = 0;
 const getID = (prefix = "py-script") => `${prefix}-${id++}`;
 
-let bootstrap = true;
+let bootstrap = true,
+    XWorker,
+    sharedRuntime;
 const sharedPyodide = new Promise((resolve) => {
     const pyConfig = document.querySelector("py-config");
     const config = pyConfig?.getAttribute("src") || pyConfig?.textContent;
     registerPlugin("py-script", {
         config,
         type: "pyodide", // or just 'py'
-        async onRuntimeReady(_, pyodide) {
+        codeBeforeRunWorker: `print('codeBeforeRunWorker')`,
+        codeAfterRunWorker: `print('codeAfterRunWorker')`,
+        onBeforeRun(pyodide, node) {
+            pyodide.runtime.globals.set("XWorker", XWorker);
+            console.log("onBeforeRun", sharedRuntime === pyodide, node);
+        },
+        onAfterRun(pyodide, node) {
+            console.log("onAfterRun", sharedRuntime === pyodide, node);
+        },
+        async onRuntimeReady(pyodide) {
             // bootstrap the shared runtime once
             // as each node as plugin gets onRuntimeReady called once
             // because no custom-element is strictly needed
             if (bootstrap) {
                 bootstrap = false;
+                sharedRuntime = pyodide;
+                XWorker = pyodide.XWorker;
                 pyodide.io.stdout = (message) => {
                     console.log("🐍", pyodide.type, message);
                 };