Support loadPyodide options (#103)

* add loadPyodideOptions * bump upstream * resolve yarn * fix loadPyodideOptions in examples * linting --------- Co-authored-by: Jeremy Tuloup <[email protected]>
jupyterlite · Mar 26, 2024 · d9b41b8 · d9b41b8
1 parent 6c57f7a
commit d9b41b8
Show file tree

Hide file tree

Showing 7 changed files with 88 additions and 23 deletions.
diff --git a/examples/jupyter-lite.json b/examples/jupyter-lite.json
@@ -2,8 +2,12 @@
   "jupyter-lite-schema-version": 0,
   "jupyter-config-data": {
     "appName": "JupyterLite Pyodide Kernel",
-    "disabledExtensions": {
-      "@jupyterlite/javascript-kernel-extension": true
+    "litePluginSettings": {
+      "@jupyterlite/pyodide-kernel-extension:kernel": {
+        "loadPyodideOptions": {
+          "packages": ["matplotlib", "micropip", "numpy", "sqlite3", "ssl"]
+        }
+      }
     }
   }
 }
diff --git a/package.json b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "@jupyterlite/pyodide-kernel-root",
-  "version": "0.2.0",
+  "version": "0.3.0",
   "private": true,
   "workspaces": {
     "packages": [
@@ -33,7 +33,7 @@
     "lint:py": "jlpm lint:py:ruff:fix",
     "lint:py:check": "jlpm lint:py:pip && jlpm lint:py:ruff:check",
     "lint:py:pip": "python -m pip check",
-    "lint:py:ruff:fix": "ruff format && ruff --fix-only",
+    "lint:py:ruff:fix": "ruff format && ruff check --fix-only",
     "lint:py:ruff:check": "ruff format --check",
     "prettier:base": "prettier \"**/*{.ts,.tsx,.js,.jsx,.css,.json,.md,.yml,.yaml}\"",
     "prettier:check": "jlpm prettier:base --check",

diff --git a/packages/pyodide-kernel-extension/schema/kernel.v0.schema.json b/packages/pyodide-kernel-extension/schema/kernel.v0.schema.json
@@ -24,6 +24,19 @@
       },
       "default": [],
       "format": "uri"
+    },
+    "loadPyodideOptions": {
+      "type": "object",
+      "description": "additional options to provide to `loadPyodide`, see https://pyodide.org/en/stable/usage/api/js-api.html#globalThis.loadPyodide",
+      "default": {},
+      "properties": {
+        "packages": {
+          "type": "array",
+          "items": {
+            "type": "string"
+          }
+        }
+      }
     }
   }
 }
diff --git a/packages/pyodide-kernel-extension/src/index.ts b/packages/pyodide-kernel-extension/src/index.ts
@@ -44,13 +44,21 @@ const kernel: JupyterLiteServerPlugin<void> = {
     const config =
       JSON.parse(PageConfig.getOption('litePluginSettings') || '{}')[PLUGIN_ID] || {};
     const url = config.pyodideUrl || PYODIDE_CDN_URL;
+
     const pyodideUrl = URLExt.parse(url).href;
     const pipliteWheelUrl = config.pipliteWheelUrl
       ? URLExt.parse(config.pipliteWheelUrl).href
       : undefined;
     const rawPipUrls = config.pipliteUrls || [];
     const pipliteUrls = rawPipUrls.map((pipUrl: string) => URLExt.parse(pipUrl).href);
     const disablePyPIFallback = !!config.disablePyPIFallback;
+    const loadPyodideOptions = config.loadPyodideOptions || {};
+
+    for (const [key, value] of Object.entries(loadPyodideOptions)) {
+      if (key.endsWith('URL') && typeof value === 'string') {
+        loadPyodideOptions[key] = URLExt.parse(key).href;
+      }
+    }
 
     kernelspecs.register({
       spec: {
@@ -81,6 +89,7 @@ const kernel: JupyterLiteServerPlugin<void> = {
           pipliteUrls,
           disablePyPIFallback,
           mountDrive,
+          loadPyodideOptions,
         });
       },
     });

diff --git a/packages/pyodide-kernel/src/kernel.ts b/packages/pyodide-kernel/src/kernel.ts
@@ -54,6 +54,7 @@ export class PyodideKernel extends BaseKernel implements IKernel {
     const { pyodideUrl } = options;
     const indexUrl = pyodideUrl.slice(0, pyodideUrl.lastIndexOf('/') + 1);
     const baseUrl = PageConfig.getBaseUrl();
+
     const pipliteUrls = [...(options.pipliteUrls || []), allJSONUrl.default];
 
     const disablePyPIFallback = !!options.disablePyPIFallback;
@@ -67,6 +68,7 @@ export class PyodideKernel extends BaseKernel implements IKernel {
       disablePyPIFallback,
       location: this.location,
       mountDrive: options.mountDrive,
+      loadPyodideOptions: options.loadPyodideOptions || {},
     };
   }
 
@@ -323,5 +325,14 @@ export namespace PyodideKernel {
      * Whether or not to mount the Emscripten drive
      */
     mountDrive: boolean;
+
+    /**
+     * additional options to provide to `loadPyodide`
+     * @see https://pyodide.org/en/stable/usage/api/js-api.html#globalThis.loadPyodide
+     */
+    loadPyodideOptions: Record<string, any> & {
+      pyodideLockURL: string;
+      packages: string[];
+    };
   }
 }
diff --git a/packages/pyodide-kernel/src/tokens.ts b/packages/pyodide-kernel/src/tokens.ts
@@ -71,5 +71,14 @@ export namespace IPyodideWorkerKernel {
      * Whether or not to mount the Emscripten drive
      */
     mountDrive: boolean;
+
+    /**
+     * additional options to provide to `loadPyodide`
+     * @see https://pyodide.org/en/stable/usage/api/js-api.html#globalThis.loadPyodide
+     */
+    loadPyodideOptions: Record<string, any> & {
+      pyodideLockURL: string;
+      packages: string[];
+    };
   }
 }
diff --git a/packages/pyodide-kernel/src/worker.ts b/packages/pyodide-kernel/src/worker.ts
@@ -50,7 +50,10 @@ export class PyodideRemoteKernel {
       importScripts(pyodideUrl);
       loadPyodide = (self as any).loadPyodide;
     }
-    this._pyodide = await loadPyodide({ indexURL: indexUrl });
+    this._pyodide = await loadPyodide({
+      indexURL: indexUrl,
+      ...options.loadPyodideOptions,
+    });
   }
 
   protected async initPackageManager(
@@ -60,38 +63,54 @@ export class PyodideRemoteKernel {
       throw new Error('Uninitialized');
     }
 
-    const { pipliteWheelUrl, disablePyPIFallback, pipliteUrls } = this._options;
+    const { pipliteWheelUrl, disablePyPIFallback, pipliteUrls, loadPyodideOptions } =
+      this._options;
 
-    await this._pyodide.loadPackage(['micropip']);
+    const preloaded = (loadPyodideOptions || {}).packages || [];
 
-    // get piplite early enough to impact pyodide dependencies
-    await this._pyodide.runPythonAsync(`
+    if (!preloaded.includes('micropip')) {
+      await this._pyodide.loadPackage(['micropip']);
+    }
+
+    if (!preloaded.includes('piplite')) {
+      await this._pyodide.runPythonAsync(`
       import micropip
       await micropip.install('${pipliteWheelUrl}', keep_going=True)
+    `);
+    }
+
+    // get piplite early enough to impact pyodide-kernel dependencies
+    await this._pyodide.runPythonAsync(`
       import piplite.piplite
       piplite.piplite._PIPLITE_DISABLE_PYPI = ${disablePyPIFallback ? 'True' : 'False'}
       piplite.piplite._PIPLITE_URLS = ${JSON.stringify(pipliteUrls)}
     `);
   }
 
   protected async initKernel(options: IPyodideWorkerKernel.IOptions): Promise<void> {
-    // from this point forward, only use piplite (but not %pip)
-    await this._pyodide.runPythonAsync(`
-      await piplite.install(['ssl'], keep_going=True);
-      await piplite.install(['sqlite3'], keep_going=True);
-      await piplite.install(['ipykernel'], keep_going=True);
-      await piplite.install(['comm'], keep_going=True);
-      await piplite.install(['pyodide_kernel'], keep_going=True);
-      await piplite.install(['ipython'], keep_going=True);
-      import pyodide_kernel
-    `);
+    const preloaded = (options.loadPyodideOptions || {}).packages || [];
+
+    const toLoad = ['ssl', 'sqlite3', 'ipykernel', 'comm', 'pyodide_kernel', 'ipython'];
+
+    const scriptLines: string[] = [];
+
+    // use piplite for packages that weren't pre-loaded
+    for (const pkgName of toLoad) {
+      if (!preloaded.includes(pkgName)) {
+        scriptLines.push(`await piplite.install('${pkgName}', keep_going=True)`);
+      }
+    }
+
+    // import the kernel
+    scriptLines.push('import pyodide_kernel');
+
     // cd to the kernel location
     if (options.mountDrive && this._localPath) {
-      await this._pyodide.runPythonAsync(`
-        import os;
-        os.chdir("${this._localPath}");
-      `);
+      scriptLines.push('import os', `os.chdir("${this._localPath}")`);
     }
+
+    // from this point forward, only use piplite (but not %pip)
+    await this._pyodide.runPythonAsync(scriptLines.join('\n'));
   }
 
   protected async initGlobals(options: IPyodideWorkerKernel.IOptions): Promise<void> {