ENH Support alternative index urls (#74)

CHANGELOG.md

@@ -10,7 +10,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Added `verbose` parameter to micropip.install and micropip.uninstall
   [#60](https://github.com/pyodide/micropip/pull/60)
-
+- Added `index_urls` parameter to micropip.install to support installing
+  from custom package indexes.
+  [#74](https://github.com/pyodide/micropip/pull/74)
+- Added `micropip.set_index_urls` to support installing from custom package
+  indexes.
+  [#74](https://github.com/pyodide/micropip/pull/74)
 ### Fixed
 
 - Fixed `micropip.add_mock_package` to work with Pyodide>=0.23.0

micropip/__init__.py

@@ -1,4 +1,5 @@
 from ._commands.freeze import freeze
+from ._commands.index_urls import set_index_urls
 from ._commands.install import install
 from ._commands.list import _list as list
 from ._commands.mock_package import (
@@ -21,5 +22,6 @@
     "list_mock_packages",
     "remove_mock_package",
     "uninstall",
+    "set_index_urls",
     "__version__",
 ]

micropip/_commands/index_urls.py

@@ -0,0 +1,27 @@
+from .. import package_index
+
+
+def set_index_urls(urls: list[str] | str) -> None:
+    """
+    Set the index URLs to use when looking up packages.
+
+    - The index URL should support the
+        `JSON API <https://warehouse.pypa.io/api-reference/json/>`__ .
+
+    - The index URL may contain the placeholder {package_name} which will be
+        replaced with the package name when looking up a package. If it does not
+        contain the placeholder, the package name will be appended to the URL.
+
+    - If a list of URLs is provided, micropip will try each URL in order until
+        it finds a package. If no package is found, an error will be raised.
+
+    Parameters
+    ----------
+    urls
+        A list of URLs or a single URL to use as the package index.
+    """
+
+    if isinstance(urls, str):
+        urls = [urls]
+
+    package_index.INDEX_URLS = urls

micropip/_commands/install.py

@@ -16,6 +16,7 @@ async def install(
     deps: bool = True,
     credentials: str | None = None,
     pre: bool = False,
+    index_urls: list[str] | str | None = None,
     *,
     verbose: bool | int = False,
 ) -> None:
@@ -86,6 +87,21 @@ async def install(
         If ``True``, include pre-release and development versions. By default,
         micropip only finds stable versions.
 
+    index_urls :
+
+        A list of URLs or a single URL to use as the package index when looking
+        up packages. If None, *https://pypi.org/pypi/{package_name}/json* is used.
+
+        - The index URL should support the
+          `JSON API <https://warehouse.pypa.io/api-reference/json/>`__ .
+
+        - The index URL may contain the placeholder {package_name} which will be
+          replaced with the package name when looking up a package. If it does not
+          contain the placeholder, the package name will be appended to the URL.
+
+        - If a list of URLs is provided, micropip will try each URL in order until
+          it finds a package. If no package is found, an error will be raised.
+
     verbose :
         Print more information about the process.
         By default, micropip is silent. Setting ``verbose=True`` will print
@@ -117,6 +133,7 @@ async def install(
         pre=pre,
         fetch_kwargs=fetch_kwargs,
         verbose=verbose,
+        index_urls=index_urls,
     )
     await transaction.gather_requirements(requirements)
 

micropip/package_index.py

@@ -1,3 +1,5 @@
+import json
+import string
 import sys
 from collections import defaultdict
 from collections.abc import Generator
@@ -7,8 +9,14 @@
 from packaging.utils import InvalidWheelFilename
 from packaging.version import InvalidVersion, Version
 
+from ._compat import fetch_string
 from ._utils import is_package_compatible, parse_version
 
+DEFAULT_INDEX_URLS = ["https://pypi.org/pypi/{package_name}/json"]
+INDEX_URLS = DEFAULT_INDEX_URLS
+
+_formatter = string.Formatter()
+
 
 # TODO: Merge this class with WheelInfo
 @dataclass
@@ -182,3 +190,60 @@ def _fast_check_incompatibility(filename: str) -> bool:
         return False
 
     return True
+
+
+def _contain_placeholder(url: str, placeholder: str = "package_name") -> bool:
+    fields = [parsed[1] for parsed in _formatter.parse(url)]
+
+    return placeholder in fields
+
+
+async def query_package(
+    name: str,
+    fetch_kwargs: dict[str, str] | None = None,
+    index_urls: list[str] | str | None = None,
+) -> ProjectInfo:
+    """
+    Query for a package from package indexes.
+
+    Parameters
+    ----------
+    name
+        Name of the package to search for.
+    fetch_kwargs
+        Keyword arguments to pass to the fetch function.
+    index_urls
+        A list of URLs or a single URL to use as the package index.
+        If None, the default index URL is used.
+
+        If a list of URLs is provided, it will be tried in order until
+        it finds a package. If no package is found, an error will be raised.
+    """
+    global INDEX_URLS
+
+    if not fetch_kwargs:
+        fetch_kwargs = {}
+
+    if index_urls is None:
+        index_urls = INDEX_URLS
+    elif isinstance(index_urls, str):
+        index_urls = [index_urls]
+
+    for url in index_urls:
+        if _contain_placeholder(url):
+            url = url.format(package_name=name)
+        else:
+            url = f"{url}/{name}/"
+
+        try:
+            metadata = await fetch_string(url, fetch_kwargs)
+        except OSError:
+            continue
+
+        return ProjectInfo.from_json_api(json.loads(metadata))
+    else:
+        raise ValueError(
+            f"Can't fetch metadata for '{name}'."
+            "Please make sure you have entered a correct package name "
+            "and correctly specified index_urls (if you changed them)."
+        )

micropip/transaction.py

@@ -16,10 +16,10 @@
 from packaging.utils import canonicalize_name
 from packaging.version import Version
 
+from . import package_index
 from ._compat import (
     REPODATA_PACKAGES,
     fetch_bytes,
-    fetch_string,
     get_dynlibs,
     loadDynlib,
     loadedPackages,
@@ -177,6 +177,7 @@ class Transaction:
     deps: bool
     pre: bool
     fetch_kwargs: dict[str, str]
+    index_urls: list[str] | str | None
 
     locked: dict[str, PackageMetadata] = field(default_factory=dict)
     wheels: list[WheelInfo] = field(default_factory=list)
@@ -295,7 +296,9 @@ def eval_marker(e: dict[str, str]) -> bool:
             )
             return
 
-        metadata: ProjectInfo = await _get_pypi_json(req.name, self.fetch_kwargs)
+        metadata = await package_index.query_package(
+            req.name, self.fetch_kwargs, index_urls=self.index_urls
+        )
 
         try:
             wheel = find_wheel(metadata, req)
@@ -405,18 +408,6 @@ def find_wheel(metadata: ProjectInfo, req: Requirement) -> WheelInfo:
     )
 
 
-async def _get_pypi_json(pkgname: str, fetch_kwargs: dict[str, str]) -> ProjectInfo:
-    url = f"https://pypi.org/pypi/{pkgname}/json"
-    try:
-        metadata = await fetch_string(url, fetch_kwargs)
-    except OSError as e:
-        raise ValueError(
-            f"Can't fetch metadata for '{pkgname}' from PyPI. "
-            "Please make sure you have entered a correct package name."
-        ) from e
-    return ProjectInfo.from_json_api(json.loads(metadata))
-
-
 def _generate_package_hash(data: IO[bytes]) -> str:
     sha256_hash = hashlib.sha256()
     data.seek(0)

pyproject.toml

@@ -16,6 +16,7 @@ dynamic = ["version"]
 dependencies = ["packaging>=23.0"]
 [project.optional-dependencies]
 test = [
+  "pytest-httpserver",
   "pytest-pyodide",
   "pytest-cov",
   "build",

tests/conftest.py

@@ -1,3 +1,5 @@
+import functools
+import gzip
 import io
 import sys
 import zipfile
@@ -14,6 +16,12 @@
 PLATFORM = f"emscripten_{EMSCRIPTEN_VER.replace('.', '_')}_wasm32"
 CPVER = f"cp{sys.version_info.major}{sys.version_info.minor}"
 
+TEST_PYPI_RESPONSE_DIR = Path(__file__).parent / "test_data" / "pypi_response"
+
+
+def _read_pypi_response(file: Path) -> bytes:
+    return gzip.decompress(file.read_bytes())
+
 
 def _build(build_dir, dist_dir):
     import build
@@ -186,7 +194,7 @@ def add_pkg_version(
         self.metadata_map[filename] = metadata
         self.top_level_map[filename] = top_level
 
-    async def _get_pypi_json(self, pkgname, kwargs):
+    async def query_package(self, pkgname, kwargs, index_urls=None):
         from micropip.package_index import ProjectInfo
 
         try:
@@ -229,9 +237,43 @@ def write_file(filename, contents):
 @pytest.fixture
 def mock_fetch(monkeypatch, mock_importlib):
     pytest.importorskip("packaging")
-    from micropip import transaction
+    from micropip import package_index, transaction
 
     result = mock_fetch_cls()
-    monkeypatch.setattr(transaction, "_get_pypi_json", result._get_pypi_json)
+    monkeypatch.setattr(package_index, "query_package", result.query_package)
     monkeypatch.setattr(transaction, "fetch_bytes", result._fetch_bytes)
     return result
+
+
+def _mock_package_index_gen(
+    httpserver,
+    pkgs=("black", "pytest", "numpy", "pytz", "snowballstemmer"),
+    content_type="application/json",
+    suffix="_json.json.gz",
+):
+    # Run a mock server that serves as a package index
+    import secrets
+
+    base = secrets.token_hex(16)
+
+    for pkg in pkgs:
+        data = _read_pypi_response(TEST_PYPI_RESPONSE_DIR / f"{pkg}{suffix}")
+        httpserver.expect_request(f"/{base}/{pkg}/").respond_with_data(
+            data,
+            content_type=content_type,
+            headers={"Access-Control-Allow-Origin": "*"},
+        )
+
+    index_url = httpserver.url_for(base)
+
+    return index_url
+
+
+@pytest.fixture
+def mock_package_index_json_api(httpserver):
+    return functools.partial(
+        _mock_package_index_gen,
+        httpserver=httpserver,
+        suffix="_json.json.gz",
+        content_type="application/json",
+    )

tests/test_install.py

@@ -368,3 +368,34 @@ async def run_test(selenium, url, name, version):
                 assert f"Successfully installed {name}-{version}" in captured
 
         run_test(selenium_standalone_micropip, wheel_url, name, version)
+
+
+@pytest.mark.asyncio
+async def test_custom_index_urls(mock_package_index_json_api, monkeypatch):
+    from io import BytesIO
+
+    mock_server_fake_package = mock_package_index_json_api(
+        pkgs=["fake-pkg-micropip-test"]
+    )
+
+    _wheel_url = ""
+
+    async def _mock_fetch_bytes(url, *args):
+        nonlocal _wheel_url
+        _wheel_url = url
+        return BytesIO(b"fake wheel")
+
+    from micropip import transaction
+
+    monkeypatch.setattr(transaction, "fetch_bytes", _mock_fetch_bytes)
+
+    try:
+        await micropip.install(
+            "fake-pkg-micropip-test", index_urls=[mock_server_fake_package]
+        )
+    except Exception:
+        # We just check that the custom index url was used
+        # install will fail because the package is not real, but it doesn't matter.
+        pass
+
+    assert "fake_pkg_micropip_test-1.0.0-py2.py3-none-any.whl" in _wheel_url