feat: Add middlewares to Client in py-rattler (#915)

Co-authored-by: Bas Zalmstra <bas@prefix.dev>

py-rattler/docs/authenticated_client.md → py-rattler/docs/client.md

@@ -1,3 +1,5 @@
-# AuthenticatedClient
+# Client
+
+::: rattler.networking
 
 ::: rattler.networking.fetch_repo_data

py-rattler/examples/install_to_prefix.py

@@ -0,0 +1,65 @@
+#!/usr/bin/env -S pixi exec --spec py-rattler --spec typer -- python
+
+import asyncio
+from pathlib import Path
+from typing import get_args
+
+from rattler import install as rattler_install
+from rattler import LockFile, Platform
+from rattler.platform.platform import PlatformLiteral
+from rattler.networking import Client, MirrorMiddleware, AuthenticationMiddleware
+import typer
+
+
+app = typer.Typer()
+
+
+async def _install(
+    lock_file_path: Path,
+    environment_name: str,
+    platform: Platform,
+    target_prefix: Path,
+) -> None:
+    lock_file = LockFile.from_path(lock_file_path)
+    environment = lock_file.environment(environment_name)
+    if environment is None:
+        raise ValueError(f"Environment {environment_name} not found in lock file {lock_file_path}")
+    records = environment.conda_repodata_records_for_platform(platform)
+    if not records:
+        raise ValueError(f"No records found for platform {platform} in lock file {lock_file_path}")
+    await rattler_install(
+        records=records,
+        target_prefix=target_prefix,
+        client=Client(
+            middlewares=[
+                MirrorMiddleware({"https://conda.anaconda.org/conda-forge": ["https://repo.prefix.dev/conda-forge"]}),
+                AuthenticationMiddleware(),
+            ]
+        ),
+    )
+
+
+@app.command()
+def install(
+    lock_file_path: Path = Path("pixi.lock").absolute(),
+    environment_name: str = "default",
+    platform: str = str(Platform.current()),
+    target_prefix: Path = Path("env").absolute(),
+) -> None:
+    """
+    Installs a pixi.lock file to a custom prefix.
+    """
+    if platform not in get_args(PlatformLiteral):
+        raise ValueError(f"Invalid platform {platform}. Must be one of {get_args(PlatformLiteral)}")
+    asyncio.run(
+        _install(
+            lock_file_path=lock_file_path,
+            environment_name=environment_name,
+            platform=Platform(platform),  # type: ignore[arg-type]
+            target_prefix=target_prefix,
+        )
+    )
+
+
+if __name__ == "__main__":
+    app()

py-rattler/mkdocs.yml

@@ -118,7 +118,7 @@ nav:
           - MatchSpec: match_spec.md
           - NamelessMatchSpec: nameless_match_spec.md
       - networking:
-          - AuthenticatedClient: authenticated_client.md
+          - Client: client.md
       - package:
           - PackageName: package_name.md
       - platform:

py-rattler/pixi.toml

@@ -34,12 +34,15 @@ pytest = "~=7.4.0"
 pytest-asyncio = "0.21.1.*"
 pytest-xprocess = ">=0.23.0,<0.24"
 
+# used in examples
+typer = "*"
+
 [feature.test.pypi-dependencies]
 types-networkx = "*"
 
 [feature.test.tasks]
 test = { cmd = "pytest --doctest-modules", depends_on = ["build"] }
-fmt-python = "ruff format rattler examples"
+fmt-python = "ruff format rattler examples tests"
 fmt-rust = "cargo fmt --all"
 lint-python = "ruff check ."
 lint-rust = "cargo clippy --all"
@@ -49,7 +52,7 @@ type-check = { cmd = "mypy", depends_on = ["build"] }
 
 # checks for the CI
 fmt-rust-check = "cargo fmt --all --check"
-fmt-python-check = "ruff format rattler examples --diff"
+fmt-python-check = "ruff format rattler examples tests --diff"
 fmt-check = { depends_on = ["fmt-python-check", "fmt-rust-check"] }
 
 [feature.docs.dependencies]

py-rattler/rattler/__init__.py

@@ -10,7 +10,7 @@
     SourceConfig,
 )
 from rattler.channel import Channel, ChannelConfig, ChannelPriority
-from rattler.networking import AuthenticatedClient, fetch_repo_data
+from rattler.networking import Client, fetch_repo_data
 from rattler.virtual_package import GenericVirtualPackage, VirtualPackage, VirtualPackageOverrides, Override
 from rattler.package import (
     PackageName,
@@ -51,7 +51,7 @@
     "Channel",
     "ChannelConfig",
     "ChannelPriority",
-    "AuthenticatedClient",
+    "Client",
     "PatchInstructions",
     "RepoDataRecord",
     "RepoData",

py-rattler/rattler/install/installer.py

@@ -2,7 +2,7 @@
 import os
 from typing import List, Optional
 
-from rattler.networking.authenticated_client import AuthenticatedClient
+from rattler.networking.client import Client
 from rattler.platform.platform import Platform
 from rattler.prefix.prefix_record import PrefixRecord
 from rattler.repo_data.record import RepoDataRecord
@@ -18,7 +18,7 @@ async def install(
     platform: Optional[Platform] = None,
     execute_link_scripts: bool = False,
     show_progress: bool = True,
-    client: Optional[AuthenticatedClient] = None,
+    client: Optional[Client] = None,
 ) -> None:
     """
     Create an environment by downloading and linking the `dependencies` in

py-rattler/rattler/networking/__init__.py

@@ -1,4 +1,5 @@
-from rattler.networking.authenticated_client import AuthenticatedClient
+from rattler.networking.client import Client
+from rattler.networking.middleware import MirrorMiddleware, AuthenticationMiddleware
 from rattler.networking.fetch_repo_data import fetch_repo_data
 
-__all__ = ["AuthenticatedClient", "fetch_repo_data"]
+__all__ = ["fetch_repo_data", "Client", "MirrorMiddleware", "AuthenticationMiddleware"]

py-rattler/rattler/networking/client.py

@@ -0,0 +1,51 @@
+from __future__ import annotations
+from rattler.rattler import PyClientWithMiddleware
+from rattler.networking.middleware import AuthenticationMiddleware, MirrorMiddleware
+
+
+class Client:
+    """
+    A client that can be used to make requests.
+    """
+
+    def __init__(self, middlewares: list[AuthenticationMiddleware | MirrorMiddleware] | None = None) -> None:
+        self._client = PyClientWithMiddleware(
+            [middleware._middleware for middleware in middlewares] if middlewares else None
+        )
+
+    @classmethod
+    def _from_ffi_object(cls, client: PyClientWithMiddleware) -> Client:
+        """
+        Construct py-rattler Client from PyClientWithMiddleware FFI object.
+        """
+        client = cls.__new__(cls)
+        client._client = client
+        return client
+
+    def __repr__(self) -> str:
+        """
+        Returns a representation of the Client
+
+        Examples
+        --------
+        ```python
+        >>> Client()
+        Client()
+        >>>
+        ```
+        """
+        return f"{type(self).__name__}()"
+
+    @staticmethod
+    def authenticated_client() -> Client:
+        """
+        Returns an authenticated client.
+
+        Examples
+        --------
+        ```python
+        >>> Client.authenticated_client()
+        Client()
+        >>>
+        """
+        return Client([AuthenticationMiddleware()])

py-rattler/rattler/networking/fetch_repo_data.py

@@ -2,6 +2,7 @@
 from typing import Callable, List, Optional, Union, TYPE_CHECKING
 
 
+from rattler.networking.client import Client
 from rattler.rattler import py_fetch_repo_data
 from rattler.repo_data.sparse import SparseRepoData
 
@@ -17,6 +18,7 @@ async def fetch_repo_data(
     platforms: List[Platform],
     cache_path: Union[str, os.PathLike[str]],
     callback: Optional[Callable[[int, int], None]],
+    client: Optional[Client] = None,
 ) -> List[SparseRepoData]:
     """
     Returns a list of RepoData for given channels and platform.
@@ -29,6 +31,7 @@ async def fetch_repo_data(
                     be downloaded.
         callback: A `Callable[[int, int], None]` to report the download
                   progress of repo data.
+        client: A `Client` to use for fetching the repo data.
 
     Returns:
         A list of `SparseRepoData` for requested channels and platforms.
@@ -38,6 +41,7 @@ async def fetch_repo_data(
         [platform._inner for platform in platforms],
         cache_path,
         callback,
+        client,
     )
 
     return [SparseRepoData._from_py_sparse_repo_data(repo_data) for repo_data in repo_data_list]

py-rattler/rattler/networking/middleware.py

@@ -0,0 +1,63 @@
+from __future__ import annotations
+from rattler.rattler import PyMirrorMiddleware, PyAuthenticationMiddleware
+
+
+class MirrorMiddleware:
+    def __init__(self, mirrors: dict[str, list[str]]) -> None:
+        """
+        Create a new MirrorMiddleware instance.
+        The mirrors argument should be a dictionary where the keys are the
+        original mirror URLs and the values are lists of mirror URLs to
+        replace the original mirror with.
+
+        Examples
+        --------
+        ```python
+        >>> from rattler.networking import Client
+        >>> middleware = MirrorMiddleware({"https://conda.anaconda.org/conda-forge": ["https://repo.prefix.dev/conda-forge"]})
+        >>> middleware
+        MirrorMiddleware()
+        >>> Client([middleware])
+        Client()
+        >>>
+        ```
+        """
+        self._middleware = PyMirrorMiddleware(mirrors)
+
+    def __repr__(self) -> str:
+        """
+        Returns a representation of the Middleware
+
+        Examples
+        --------
+        ```python
+        >>> middleware = MirrorMiddleware({"https://conda.anaconda.org/conda-forge": ["https://repo.prefix.dev/conda-forge"]})
+        >>> middleware
+        MirrorMiddleware()
+        >>>
+        ```
+        """
+        return f"{type(self).__name__}()"
+
+
+class AuthenticationMiddleware:
+    def __init__(self) -> None:
+        self._middleware = PyAuthenticationMiddleware()
+
+    def __repr__(self) -> str:
+        """
+        Returns a representation of the Middleware
+
+        Examples
+        --------
+        ```python
+        >>> from rattler.networking import Client
+        >>> middleware = AuthenticationMiddleware()
+        >>> middleware
+        AuthenticationMiddleware()
+        >>> Client([middleware])
+        Client()
+        >>>
+        ```
+        """
+        return f"{type(self).__name__}()"

py-rattler/src/installer.rs

@@ -9,8 +9,8 @@ use rattler::{
 use rattler_conda_types::{PrefixRecord, RepoDataRecord};
 
 use crate::{
-    error::PyRattlerError, networking::authenticated_client::PyAuthenticatedClient,
-    platform::PyPlatform, record::PyRecord,
+    error::PyRattlerError, networking::client::PyClientWithMiddleware, platform::PyPlatform,
+    record::PyRecord,
 };
 
 // TODO: Accept functions to report progress
@@ -23,7 +23,7 @@ pub fn py_install<'a>(
     execute_link_scripts: bool,
     show_progress: bool,
     platform: Option<PyPlatform>,
-    client: Option<PyAuthenticatedClient>,
+    client: Option<PyClientWithMiddleware>,
     cache_dir: Option<PathBuf>,
     installed_packages: Option<Vec<&'a PyAny>>,
 ) -> PyResult<&'a PyAny> {

py-rattler/src/lib.rs

@@ -47,7 +47,8 @@ use lock::{
 use match_spec::PyMatchSpec;
 use meta::get_rattler_version;
 use nameless_match_spec::PyNamelessMatchSpec;
-use networking::{authenticated_client::PyAuthenticatedClient, py_fetch_repo_data};
+use networking::middleware::{PyAuthenticationMiddleware, PyMirrorMiddleware};
+use networking::{client::PyClientWithMiddleware, py_fetch_repo_data};
 use no_arch_type::PyNoArchType;
 use package_name::PyPackageName;
 use paths_json::{PyFileMode, PyPathType, PyPathsEntry, PyPathsJson, PyPrefixPlaceholder};
@@ -96,7 +97,9 @@ fn rattler(py: Python<'_>, m: &PyModule) -> PyResult<()> {
     m.add_class::<PyPlatform>().unwrap();
     m.add_class::<PyArch>().unwrap();
 
-    m.add_class::<PyAuthenticatedClient>().unwrap();
+    m.add_class::<PyMirrorMiddleware>().unwrap();
+    m.add_class::<PyAuthenticationMiddleware>().unwrap();
+    m.add_class::<PyClientWithMiddleware>().unwrap();
 
     // Shell activation things
     m.add_class::<PyActivationVariables>().unwrap();