feat(py): add custom progress reporter callbacks to installer

py-rattler/rattler/__init__.py

@@ -30,7 +30,7 @@
 from rattler.prefix import PrefixRecord, PrefixPaths, PrefixPathsEntry, PrefixPathType, Link, LinkType
 from rattler.platform import Platform
 from rattler.utils.rattler_version import get_rattler_version as _get_rattler_version
-from rattler.install import install
+from rattler.install import install, InstallerReporter
 from rattler.index import index
 from rattler.lock import (
     LockFile,
@@ -87,6 +87,7 @@
     "solve_with_sparse_repodata",
     "Platform",
     "install",
+    "InstallerReporter",
     "index",
     "AboutJson",
     "RunExportsJson",

py-rattler/rattler/install/__init__.py

@@ -1,3 +1,3 @@
-from rattler.install.installer import install
+from rattler.install.installer import install, InstallerReporter
 
-__all__ = ["install"]
+__all__ = ["install", "InstallerReporter"]

py-rattler/rattler/install/installer.py

@@ -11,6 +11,268 @@
 from rattler.rattler import py_install
 
 
+class InstallerReporter:
+    """
+    Base class for custom installation progress reporters.
+
+    Subclass this and override the methods we care about to receive progress
+    callbacks during package installation. All methods have no-op defaults so
+    we only need to implement the ones relevant to our use case.
+
+    Methods that return an ``int`` act as opaque tokens: the value we return
+    will be passed back to the corresponding ``*_complete`` callback so we can
+    correlate start/finish events.
+
+    Example
+    -------
+    ```python
+    from rattler.install import InstallerReporter
+
+    class MyReporter(InstallerReporter):
+        def on_transaction_start(self, total_operations: int) -> None:
+            print(f"Starting installation of {total_operations} operations")
+
+        def on_download_progress(
+            self,
+            download_idx: int,
+            progress: int,
+            total: int | None,
+        ) -> None:
+            pct = f"{progress}/{total}" if total else str(progress)
+            print(f"  downloading [{download_idx}]: {pct} bytes")
+
+        def on_transaction_complete(self) -> None:
+            print("Installation complete!")
+    ```
+    """
+
+    def on_transaction_start(self, total_operations: int) -> None:
+        """Called once when the installation transaction begins.
+
+        Parameters
+        ----------
+        total_operations:
+            The total number of install/unlink operations in this transaction.
+        """
+
+    def on_transaction_operation_start(self, operation: int) -> None:
+        """Called when a single transaction operation starts.
+
+        Parameters
+        ----------
+        operation:
+            Index of the operation within the transaction.
+        """
+
+    def on_populate_cache_start(self, operation: int, package_name: str) -> int:
+        """Called when rattler begins populating the local cache for a package.
+
+        Parameters
+        ----------
+        operation:
+            Index of the operation within the transaction.
+        package_name:
+            Normalised name of the package being cached.
+
+        Returns
+        -------
+        int
+            An opaque token passed back to :meth:`on_populate_cache_complete`.
+        """
+        return 0
+
+    def on_validate_start(self, cache_entry: int) -> int:
+        """Called when cache-entry validation begins.
+
+        Parameters
+        ----------
+        cache_entry:
+            Token returned by :meth:`on_populate_cache_start`.
+
+        Returns
+        -------
+        int
+            An opaque token passed back to :meth:`on_validate_complete`.
+        """
+        return 0
+
+    def on_validate_complete(self, validate_idx: int) -> None:
+        """Called when cache-entry validation finishes.
+
+        Parameters
+        ----------
+        validate_idx:
+            Token returned by :meth:`on_validate_start`.
+        """
+
+    def on_download_start(self, cache_entry: int) -> int:
+        """Called just before a package download begins.
+
+        Parameters
+        ----------
+        cache_entry:
+            Token returned by :meth:`on_populate_cache_start`.
+
+        Returns
+        -------
+        int
+            An opaque token passed to :meth:`on_download_progress` and
+            :meth:`on_download_completed`.
+        """
+        return 0
+
+    def on_download_progress(self, download_idx: int, progress: int, total: Optional[int]) -> None:
+        """Called periodically with download byte progress.
+
+        Parameters
+        ----------
+        download_idx:
+            Token returned by :meth:`on_download_start`.
+        progress:
+            Bytes downloaded so far.
+        total:
+            Total expected bytes, or ``None`` if unknown.
+        """
+
+    def on_download_completed(self, download_idx: int) -> None:
+        """Called when a download finishes.
+
+        Parameters
+        ----------
+        download_idx:
+            Token returned by :meth:`on_download_start`.
+        """
+
+    def on_populate_cache_complete(self, cache_entry: int) -> None:
+        """Called when the cache has been fully populated for a package.
+
+        Parameters
+        ----------
+        cache_entry:
+            Token returned by :meth:`on_populate_cache_start`.
+        """
+
+    def on_unlink_start(self, operation: int, package_name: str) -> int:
+        """Called when rattler begins unlinking (removing) a package.
+
+        Parameters
+        ----------
+        operation:
+            Index of the operation within the transaction.
+        package_name:
+            Normalised name of the package being unlinked.
+
+        Returns
+        -------
+        int
+            An opaque token passed back to :meth:`on_unlink_complete`.
+        """
+        return 0
+
+    def on_unlink_complete(self, index: int) -> None:
+        """Called when a package has been fully unlinked.
+
+        Parameters
+        ----------
+        index:
+            Token returned by :meth:`on_unlink_start`.
+        """
+
+    def on_link_start(self, operation: int, package_name: str) -> int:
+        """Called when rattler begins linking (installing) a package.
+
+        Parameters
+        ----------
+        operation:
+            Index of the operation within the transaction.
+        package_name:
+            Normalised name of the package being linked.
+
+        Returns
+        -------
+        int
+            An opaque token passed back to :meth:`on_link_complete`.
+        """
+        return 0
+
+    def on_link_complete(self, index: int) -> None:
+        """Called when a package has been fully linked into the prefix.
+
+        Parameters
+        ----------
+        index:
+            Token returned by :meth:`on_link_start`.
+        """
+
+    def on_transaction_operation_complete(self, operation: int) -> None:
+        """Called when a single transaction operation finishes.
+
+        Parameters
+        ----------
+        operation:
+            Index of the operation within the transaction.
+        """
+
+    def on_transaction_complete(self) -> None:
+        """Called once when the entire transaction has finished."""
+
+    def on_post_link_start(self, package_name: str, script_path: str) -> int:
+        """Called when a post-link script begins execution.
+
+        Parameters
+        ----------
+        package_name:
+            Name of the package whose post-link script is running.
+        script_path:
+            Relative path to the script within the prefix.
+
+        Returns
+        -------
+        int
+            An opaque token passed back to :meth:`on_post_link_complete`.
+        """
+        return 0
+
+    def on_post_link_complete(self, index: int, success: bool) -> None:
+        """Called when a post-link script finishes.
+
+        Parameters
+        ----------
+        index:
+            Token returned by :meth:`on_post_link_start`.
+        success:
+            ``True`` if the script exited successfully.
+        """
+
+    def on_pre_unlink_start(self, package_name: str, script_path: str) -> int:
+        """Called when a pre-unlink script begins execution.
+
+        Parameters
+        ----------
+        package_name:
+            Name of the package whose pre-unlink script is running.
+        script_path:
+            Relative path to the script within the prefix.
+
+        Returns
+        -------
+        int
+            An opaque token passed back to :meth:`on_pre_unlink_complete`.
+        """
+        return 0
+
+    def on_pre_unlink_complete(self, index: int, success: bool) -> None:
+        """Called when a pre-unlink script finishes.
+
+        Parameters
+        ----------
+        index:
+            Token returned by :meth:`on_pre_unlink_start`.
+        success:
+            ``True`` if the script exited successfully.
+        """
+
+
 async def install(
     records: List[RepoDataRecord],
     target_prefix: str | os.PathLike[str],
@@ -23,6 +285,7 @@ async def install(
     show_progress: bool = True,
     client: Optional[Client] = None,
     requested_specs: Optional[List[MatchSpec]] = None,
+    reporter: Optional[InstallerReporter] = None,
 ) -> None:
     """
     Create an environment by downloading and linking the `dependencies` in
@@ -74,11 +337,15 @@ async def install(
         execute_link_scripts: whether to execute the post-link and pre-unlink scripts
                 that may be part of a package. Defaults to False.
         show_progress: If set to `True` a progress bar will be shown on the CLI.
+                Ignored when `reporter` is provided.
         client: An authenticated client to use for downloading packages. If not specified a default
                 client will be used.
         requested_specs: A list of `MatchSpec`s that were originally requested. These will be used
                 to populate the `requested_specs` field in the generated `conda-meta/*.json` files.
                 If `None`, the `requested_specs` field will remain empty.
+        reporter: An optional :class:`InstallerReporter` instance that receives progress
+                callbacks during installation. When provided, `show_progress` is ignored.
+                Subclass :class:`InstallerReporter` and override the methods you need.
     """
 
     await py_install(
@@ -93,4 +360,5 @@ async def install(
         execute_link_scripts=execute_link_scripts,
         show_progress=show_progress,
         requested_specs=requested_specs,
+        reporter=reporter,
     )