ENH Add loadPyodide packages option for loading packages during b…

…ootstrap (#4100)

For improved loading performance.

conftest.py

@@ -33,7 +33,6 @@
     pyodide._api.importlib.invalidate_caches;
     pyodide._api.package_loader.unpack_buffer;
     pyodide._api.package_loader.get_dynlibs;
-    pyodide._api.package_loader.sub_resource_hash;
     pyodide.runPython("");
     pyodide.pyimport("pyodide.ffi.wrappers").destroy();
     pyodide.pyimport("pyodide.http").destroy();

docs/project/changelog.md

@@ -16,6 +16,11 @@ myst:
 
 ## Unreleased
 
+- {{ Performance }} Added a `packages` optional argument to `loadPyodide`.
+  Passing packages here saves time by downloading them during the Pyodide
+  bootstrap.
+  {pr}`4100`
+
 - {{ Performance }} Improved performance of PyProxy creation.
   {pr}`4096`
 

src/js/compat.ts

@@ -234,3 +234,23 @@ async function nodeLoadScript(url: string) {
     await import(/* webpackIgnore: true */ nodeUrlMod.pathToFileURL(url).href);
   }
 }
+
+// consider dropping this this once we drop support for node 14?
+function nodeBase16ToBase64(b16: string): string {
+  return Buffer.from(b16, "hex").toString("base64");
+}
+
+function browserBase16ToBase64(b16: string): string {
+  return btoa(
+    b16
+      .match(/\w{2}/g)!
+      .map(function (a) {
+        return String.fromCharCode(parseInt(a, 16));
+      })
+      .join(""),
+  );
+}
+
+export const base16ToBase64 = IN_NODE
+  ? nodeBase16ToBase64
+  : browserBase16ToBase64;

src/js/load-package.ts

@@ -6,6 +6,7 @@ import {
   loadBinaryFile,
   initNodeModules,
   resolvePath,
+  base16ToBase64,
 } from "./compat.js";
 import { createLock } from "./lock";
 import { loadDynlibsFromPackage } from "./dynload";
@@ -60,6 +61,7 @@ async function initializePackageIndex(lockFileURL: string) {
     API.lockfile_unvendored_stdlibs_and_test.filter(
       (lib: string) => lib !== "test",
     );
+  await loadPackage(API.config.packages, { messageCallback() {} });
 }
 
 API.packageIndexReady = initializePackageIndex(API.config.lockFileURL);
@@ -244,9 +246,8 @@ async function downloadPackage(
     }
     file_name = API.lockfile_packages[name].file_name;
     uri = resolvePath(file_name, installBaseUrl);
-    file_sub_resource_hash = API.package_loader.sub_resource_hash(
-      API.lockfile_packages[name].sha256,
-    );
+    file_sub_resource_hash =
+      "sha256-" + base16ToBase64(API.lockfile_packages[name].sha256);
   } else {
     uri = channel;
     file_sub_resource_hash = undefined;
@@ -342,14 +343,16 @@ async function downloadAndInstall(
 
   try {
     const buffer = await downloadPackage(pkg.name, pkg.channel, checkIntegrity);
-    const installPromisDependencies = pkg.depends.map((dependency) => {
+    const installPromiseDependencies = pkg.depends.map((dependency) => {
       return toLoad.has(dependency)
         ? toLoad.get(dependency)!.done
         : Promise.resolve();
     });
+    // Can't install until bootstrap is finalized.
+    await API.bootstrapFinalizedPromise;
 
     // wait until all dependencies are installed
-    await Promise.all(installPromisDependencies);
+    await Promise.all(installPromiseDependencies);
 
     await installPackage(pkg.name, buffer, pkg.channel);
     loaded.add(pkg.name);

src/js/pyodide.ts

@@ -189,6 +189,7 @@ export type ConfigType = {
   args: string[];
   _node_mounts: string[];
   env: { [key: string]: string };
+  packages: string[];
 };
 
 /**
@@ -291,7 +292,15 @@ export async function loadPyodide(
      * `"/home/pyodide"`
      */
     env?: { [key: string]: string };
-
+    /**
+     * A list of packages to load as Pyodide is initializing.
+     *
+     * This is the same as loading the packages with
+     * :js:func:`pyodide.loadPackage` after Pyodide is loaded except using the
+     * ``packages`` option is more efficient because the packages are downloaded
+     * while Pyodide bootstraps itself.
+     */
+    packages?: string[];
     /**
      * @ignore
      */
@@ -315,6 +324,7 @@ export async function loadPyodide(
     _node_mounts: [],
     env: {},
     packageCacheDir: indexURL,
+    packages: [],
   };
   const config = Object.assign(default_config, options) as ConfigType;
   if (options.homedir) {
@@ -343,6 +353,10 @@ export async function loadPyodide(
   initializeFileSystem(Module, config);
 
   const moduleLoaded = new Promise((r) => (Module.postRun = r));
+  let bootstrapFinalized: () => void;
+  API.bootstrapFinalizedPromise = new Promise<void>(
+    (r) => (bootstrapFinalized = r),
+  );
 
   // locateFile tells Emscripten where to find the data files that initialize
   // the file system.
@@ -390,6 +404,7 @@ If you updated the Pyodide version, make sure you also updated the 'indexURL' pa
   }
 
   const pyodide = finalizeBootstrap(API, config);
+  bootstrapFinalized!();
 
   // runPython works starting here.
   if (!pyodide.version.includes("dev")) {

src/py/pyodide/_package_loader.py

@@ -1,5 +1,3 @@
-import base64
-import binascii
 import re
 import shutil
 import sys
@@ -360,23 +358,3 @@ def init_loaded_packages() -> None:
     """
     for dist in importlib_distributions():
         setattr(loadedPackages, dist.name, get_dist_source(dist))
-
-
-def sub_resource_hash(sha_256: str) -> str:
-    """Calculates the sub resource integrity hash given a SHA-256
-
-    Parameters
-    ----------
-    sha_256
-        A hexdigest of the SHA-256 sum.
-
-    Returns
-    -------
-        The sub resource integrity hash corresponding to the sum.
-
-    >>> sha_256 = 'c0dc86efda0060d4084098a90ec92b3d4aa89d7f7e0fba5424561d21451e1758'
-    >>> sub_resource_hash(sha_256)
-    'sha256-wNyG79oAYNQIQJipDskrPUqonX9+D7pUJFYdIUUeF1g='
-    """
-    binary_digest = binascii.unhexlify(sha_256)
-    return "sha256-" + base64.b64encode(binary_digest).decode()

src/tests/test_package_loading.py

@@ -515,7 +515,7 @@ def test_custom_lockfile(selenium_standalone_noload):
     selenium = selenium_standalone_noload
     lock = selenium.run_js(
         """
-        let pyodide = await loadPyodide({fullStdLib: false});
+        let pyodide = await loadPyodide({fullStdLib: false, packages: ["micropip"]});
         await pyodide.loadPackage("micropip")
         return pyodide.runPythonAsync(`
             import micropip
@@ -531,8 +531,7 @@ def test_custom_lockfile(selenium_standalone_noload):
         assert (
             selenium.run_js(
                 """
-                let pyodide = await loadPyodide({fullStdLib: false, lockFileURL: "custom_lockfile.json" });
-                await pyodide.loadPackage("hypothesis");
+                let pyodide = await loadPyodide({fullStdLib: false, lockFileURL: "custom_lockfile.json", packages: ["hypothesis"] });
                 return pyodide.runPython("import hypothesis; hypothesis.__version__")
                 """
             )