Support passing of index URLs to piplite (CLI, runtime configuration, and requirements files)

packages/pyodide-kernel-extension/schema/kernel.v0.schema.json

@@ -25,6 +25,21 @@
       "default": [],
       "format": "uri"
     },
+    "pipliteInstallDefaultOptions": {
+      "type": "object",
+      "description": "Default options to pass to piplite.install",
+      "default": {},
+      "properties": {
+        "indexUrls": {
+          "type": "array",
+          "items": {
+            "type": "string",
+            "format": "uri"
+          },
+          "description": "Base URLs of extra indices to use"
+        }
+      }
+    },
     "loadPyodideOptions": {
       "type": "object",
       "description": "additional options to provide to `loadPyodide`, see https://pyodide.org/en/stable/usage/api/js-api.html#globalThis.loadPyodide",

packages/pyodide-kernel-extension/src/index.ts

@@ -58,6 +58,10 @@ const kernel: JupyterLiteServerPlugin<void> = {
     const pipliteUrls = rawPipUrls.map((pipUrl: string) => URLExt.parse(pipUrl).href);
     const disablePyPIFallback = !!config.disablePyPIFallback;
     const loadPyodideOptions = config.loadPyodideOptions || {};
+    const pipliteInstallDefaultOptions = config.pipliteInstallDefaultOptions || {};
+
+    // Parse any configured index URLs
+    const indexUrls = pipliteInstallDefaultOptions.indexUrls || [];
 
     for (const [key, value] of Object.entries(loadPyodideOptions)) {
       if (key.endsWith('URL') && typeof value === 'string') {
@@ -99,6 +103,10 @@ const kernel: JupyterLiteServerPlugin<void> = {
           mountDrive,
           loadPyodideOptions,
           contentsManager,
+          pipliteInstallDefaultOptions: {
+            indexUrls,
+            ...pipliteInstallDefaultOptions,
+          },
         });
       },
     });

packages/pyodide-kernel/py/piplite/piplite/cli.py

@@ -13,7 +13,7 @@ async def install(
         deps: bool = True,                              # --no-deps
         credentials: str | None = None,                 # no CLI alias
         pre: bool = False,                              # --pre
-        index_urls: list[str] | str | None = None,      # no CLI alias
+        index_urls: list[str] | str | None = None,      # -i and --index-url
         *,
         verbose: bool | int | None = None,
     ):
@@ -26,10 +26,32 @@ async def install(
 import re
 import sys
 import typing
+from typing import Optional, List, Tuple
+from dataclasses import dataclass
+
 from argparse import ArgumentParser
 from pathlib import Path
 
+
+@dataclass
+class RequirementsContext:
+    """Track state while parsing requirements files."""
+
+    index_url: Optional[str] = None
+    requirements: List[str] = None
+
+    def __post_init__(self):
+        if self.requirements is None:
+            self.requirements = []
+
+    def add_requirement(self, req: str):
+        """Add a requirement with the currently active index URL."""
+        self.requirements.append((req, self.index_url))
+
+
 REQ_FILE_PREFIX = r"^(-r|--requirements)\s*=?\s*(.*)\s*"
+INDEX_URL_PREFIX = r"^(--index-url|-i)\s*=?\s*(.*)\s*"
+
 
 __all__ = ["get_transformed_code"]
 
@@ -76,6 +98,12 @@ def _get_parser() -> ArgumentParser:
         action="store_true",
         help="whether pre-release packages should be considered",
     )
+    parser.add_argument(
+        "--index-url",
+        "-i",
+        type=str,
+        help="the index URL to use for package lookup",
+    )
     parser.add_argument(
         "packages",
         nargs="*",
@@ -111,11 +139,44 @@ async def get_action_kwargs(argv: list[str]) -> tuple[typing.Optional[str], dict
         return None, {}
 
     kwargs = {}
-
     action = args.action
 
     if action == "install":
-        kwargs["requirements"] = args.packages
+        all_requirements = []
+
+        if args.packages:
+            all_requirements.extend((pkg, args.index_url) for pkg in args.packages)
+
+        # Process requirements files
+        for req_file in args.requirements or []:
+            context = RequirementsContext()
+
+            if not Path(req_file).exists():
+                warn(f"piplite could not find requirements file {req_file}")
+                continue
+
+            for line_no, line in enumerate(
+                Path(req_file).read_text(encoding="utf-8").splitlines()
+            ):
+                await _packages_from_requirements_line(
+                    Path(req_file), line_no + 1, line, context
+                )
+
+            all_requirements.extend(context.requirements)
+
+        if all_requirements:
+            kwargs["requirements"] = []
+            used_index = None
+
+            for req, idx in all_requirements:
+                if idx:
+                    used_index = idx
+                kwargs["requirements"].append(req)
+
+            # Set the index URL if one was found (either passed to the CLI or
+            # passed within the requirements file)
+            if used_index:
+                kwargs["index_urls"] = used_index
 
         if args.pre:
             kwargs["pre"] = True
@@ -126,51 +187,62 @@ async def get_action_kwargs(argv: list[str]) -> tuple[typing.Optional[str], dict
         if args.verbose:
             kwargs["keep_going"] = True
 
-        for req_file in args.requirements or []:
-            kwargs["requirements"] += await _packages_from_requirements_file(
-                Path(req_file)
-            )
-
     return action, kwargs
 
 
-async def _packages_from_requirements_file(req_path: Path) -> list[str]:
-    """Extract (potentially nested) package requirements from a requirements file."""
+async def _packages_from_requirements_file(
+    req_path: Path,
+) -> Tuple[List[str], Optional[str]]:
+    """Extract (potentially nested) package requirements from a requirements file.
+
+    Returns:
+    Tuple of (list of package requirements, optional index URL)
+    """
     if not req_path.exists():
         warn(f"piplite could not find requirements file {req_path}")
         return []
 
-    requirements = []
+    context = RequirementsContext()
 
-    for line_no, line in enumerate(req_path.read_text(encoding="utf").splitlines()):
-        requirements += await _packages_from_requirements_line(
-            req_path, line_no + 1, line
-        )
+    for line_no, line in enumerate(req_path.read_text(encoding="utf-8").splitlines()):
+        await _packages_from_requirements_line(req_path, line_no + 1, line, context)
 
-    return requirements
+    return context.requirements, context.index_url
 
 
 async def _packages_from_requirements_line(
-    req_path: Path, line_no: int, line: str
-) -> list[str]:
+    req_path: Path, line_no: int, line: str, context: RequirementsContext
+) -> None:
     """Extract (potentially nested) package requirements from line of a
     requirements file.
 
     `micropip` has a sufficient pep508 implementation to handle most cases
     """
     req = line.strip().split("#")[0].strip()
-    # is it another requirement file?
+    if not req:
+        return
+
+    # Check for nested requirements file
     req_file_match = re.match(REQ_FILE_PREFIX, req)
     if req_file_match:
-        if req_file_match[2].startswith("/"):
-            sub_req = Path(req)
+        sub_path = req_file_match[2]
+        if sub_path.startswith("/"):
+            sub_req = Path(sub_path)
         else:
-            sub_req = req_path.parent / req_file_match[2]
-        return await _packages_from_requirements_file(sub_req)
+            sub_req = req_path.parent / sub_path
+        sub_reqs, _ = await _packages_from_requirements_file(sub_req)
+        # Use current context's index_url for nested requirements
+        context.requirements.extend(sub_reqs)
+        return
+
+    # Check for index URL specification
+    index_match = re.match(INDEX_URL_PREFIX, req)
+    if index_match:
+        context.index_url = index_match[2].strip()
+        return
 
     if req.startswith("-"):
         warn(f"{req_path}:{line_no}: unrecognized requirement: {req}")
-        req = None
-    if not req:
-        return []
-    return [req]
+        return
+
+    context.add_requirement(req)

packages/pyodide-kernel/py/piplite/piplite/piplite.py

@@ -33,6 +33,9 @@
 #: a well-known file name respected by the rest of the build chain
 ALL_JSON = "/all.json"
 
+#: default index URLs to use when no specific index URLs are provided
+_PIPLITE_DEFAULT_INDEX_URLS = None
+
 
 class PiplitePyPIDisabled(ValueError):
     """An error for when PyPI is disabled at the site level, but a download was
@@ -117,17 +120,33 @@ async def _install(
     verbose: bool | int = False,
 ):
     """Invoke micropip.install with a patch to get data from local indexes"""
-    with patch("micropip.package_index.query_package", _query_package):
-        return await micropip.install(
-            requirements=requirements,
-            keep_going=keep_going,
-            deps=deps,
-            credentials=credentials,
-            pre=pre,
-            index_urls=index_urls,
-            verbose=verbose,
+
+    try:
+        # Use default index URLs if none provided and defaults exist
+        effective_index_urls = (
+            index_urls if index_urls is not None else _PIPLITE_DEFAULT_INDEX_URLS
         )
 
+        if verbose:
+            logger.info(f"Installing with index URLs: {effective_index_urls}")
+
+        with patch("micropip.package_index.query_package", _query_package):
+            return await micropip.install(
+                requirements=requirements,
+                keep_going=keep_going,
+                deps=deps,
+                credentials=credentials,
+                pre=pre,
+                index_urls=effective_index_urls,
+                verbose=verbose,
+            )
+    except Exception as e:
+        if effective_index_urls:
+            logger.error(
+                f"Failed to install using index URLs {effective_index_urls}: {e}"
+            )
+        raise
+
 
 def install(
     requirements: str | list[str],

packages/pyodide-kernel/src/kernel.ts

@@ -381,6 +381,11 @@ export namespace PyodideKernel {
      */
     mountDrive: boolean;
 
+    /**
+     * Default options to pass to piplite.install
+     */
+    pipliteInstallDefaultOptions?: IPyodideWorkerKernel.IPipliteInstallOptions;
+
     /**
      * additional options to provide to `loadPyodide`
      * @see https://pyodide.org/en/stable/usage/api/js-api.html#globalThis.loadPyodide

packages/pyodide-kernel/src/tokens.ts

@@ -48,9 +48,24 @@ export interface IPyodideWorkerKernel extends IWorkerKernel {
 export type IRemotePyodideWorkerKernel = IPyodideWorkerKernel;
 
 /**
- * An namespace for Pyodide workers.
+ * A namespace for Pyodide workers.
  */
 export namespace IPyodideWorkerKernel {
+  /**
+   * Options for piplite installation.
+   */
+  export interface IPipliteInstallOptions {
+    /**
+     * Base URLs of extra indices to use
+     */
+    indexUrls?: string[];
+
+    /**
+     * Any additional piplite install options
+     */
+    [key: string]: any;
+  }
+
   /**
    * Initialization options for a worker.
    */
@@ -90,6 +105,11 @@ export namespace IPyodideWorkerKernel {
      */
     mountDrive: boolean;
 
+    /**
+     * Default options to pass to piplite.install
+     */
+    pipliteInstallDefaultOptions?: IPyodideWorkerKernel.IPipliteInstallOptions;
+
     /**
      * additional options to provide to `loadPyodide`
      * @see https://pyodide.org/en/stable/usage/api/js-api.html#globalThis.loadPyodide

packages/pyodide-kernel/src/worker.ts

@@ -65,8 +65,13 @@ export class PyodideRemoteKernel {
       throw new Error('Uninitialized');
     }
 
-    const { pipliteWheelUrl, disablePyPIFallback, pipliteUrls, loadPyodideOptions } =
-      this._options;
+    const {
+      pipliteWheelUrl,
+      disablePyPIFallback,
+      pipliteUrls,
+      loadPyodideOptions,
+      pipliteInstallDefaultOptions,
+    } = this._options;
 
     const preloaded = (loadPyodideOptions || {}).packages || [];
 
@@ -81,12 +86,20 @@ export class PyodideRemoteKernel {
     `);
     }
 
+    const pythonConfig = [
+      'import piplite.piplite',
+      `piplite.piplite._PIPLITE_DISABLE_PYPI = ${disablePyPIFallback ? 'True' : 'False'}`,
+      `piplite.piplite._PIPLITE_URLS = ${JSON.stringify(pipliteUrls)}`,
+    ];
+
+    if (pipliteInstallDefaultOptions?.indexUrls) {
+      pythonConfig.push(
+        `piplite.piplite._PIPLITE_DEFAULT_INDEX_URLS = ${JSON.stringify(pipliteInstallDefaultOptions.indexUrls)}`,
+      );
+    }
+
     // 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)}
-    `);
+    await this._pyodide.runPythonAsync(pythonConfig.join('\n'));
   }
 
   protected async initKernel(options: IPyodideWorkerKernel.IOptions): Promise<void> {