diff --git a/py-rattler/Cargo.lock b/py-rattler/Cargo.lock index 3dc5df6df..b048c6b99 100644 --- a/py-rattler/Cargo.lock +++ b/py-rattler/Cargo.lock @@ -4244,6 +4244,7 @@ dependencies = [ "serde_json", "tempfile", "thiserror 2.0.18", + "tokio", "tracing", "url", ] diff --git a/py-rattler/rattler/__init__.py b/py-rattler/rattler/__init__.py index a95d54876..545a31dc2 100644 --- a/py-rattler/rattler/__init__.py +++ b/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", diff --git a/py-rattler/rattler/install/__init__.py b/py-rattler/rattler/install/__init__.py index e57b30395..c076b57fa 100644 --- a/py-rattler/rattler/install/__init__.py +++ b/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"] diff --git a/py-rattler/rattler/install/installer.py b/py-rattler/rattler/install/installer.py index de0f4f572..7dbf73196 100644 --- a/py-rattler/rattler/install/installer.py +++ b/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, ) diff --git a/py-rattler/src/installer.rs b/py-rattler/src/installer.rs index 616039b2d..ea0c2648d 100644 --- a/py-rattler/src/installer.rs +++ b/py-rattler/src/installer.rs @@ -1,9 +1,10 @@ use std::path::PathBuf; +use std::sync::Arc; -use pyo3::{exceptions::PyTypeError, pyfunction, Bound, PyAny, PyResult, Python}; +use pyo3::{exceptions::PyTypeError, pyfunction, Bound, PyAny, PyObject, PyResult, Python}; use pyo3_async_runtimes::tokio::future_into_py; use rattler::{ - install::{IndicatifReporter, Installer}, + install::{IndicatifReporter, Installer, Reporter, Transaction}, package_cache::PackageCache, }; use rattler_conda_types::{PackageName, PrefixRecord, RepoDataRecord}; @@ -15,10 +16,205 @@ use crate::{ record::PyRecord, }; -// TODO: Accept functions to report progress +/// A [`Reporter`] implementation that delegates progress events to a Python object. The Python object should implement the following methods: +struct PyReporter { + py_obj: Arc, +} + +impl Reporter for PyReporter { + fn on_transaction_start(&self, transaction: &Transaction) { + let total = transaction.operations.len(); + Python::with_gil(|py| { + let _ = self + .py_obj + .call_method1(py, "on_transaction_start", (total,)); + }); + } + + fn on_transaction_operation_start(&self, operation: usize) { + Python::with_gil(|py| { + let _ = self + .py_obj + .call_method1(py, "on_transaction_operation_start", (operation,)); + }); + } + + fn on_populate_cache_start(&self, operation: usize, record: &RepoDataRecord) -> usize { + let name = record.package_record.name.as_normalized().to_string(); + Python::with_gil(|py| { + match self + .py_obj + .call_method1(py, "on_populate_cache_start", (operation, name)) + { + Ok(val) => val.extract::(py).unwrap_or(0), + Err(_) => 0, + } + }) + } + + fn on_validate_start(&self, cache_entry: usize) -> usize { + Python::with_gil(|py| { + match self + .py_obj + .call_method1(py, "on_validate_start", (cache_entry,)) + { + Ok(val) => val.extract::(py).unwrap_or(0), + Err(_) => 0, + } + }) + } + + fn on_validate_complete(&self, validate_idx: usize) { + Python::with_gil(|py| { + let _ = self + .py_obj + .call_method1(py, "on_validate_complete", (validate_idx,)); + }); + } + + fn on_download_start(&self, cache_entry: usize) -> usize { + Python::with_gil(|py| { + match self + .py_obj + .call_method1(py, "on_download_start", (cache_entry,)) + { + Ok(val) => val.extract::(py).unwrap_or(0), + Err(_) => 0, + } + }) + } + + fn on_download_progress(&self, download_idx: usize, progress: u64, total: Option) { + Python::with_gil(|py| { + let _ = self.py_obj.call_method1( + py, + "on_download_progress", + (download_idx, progress, total), + ); + }); + } + + fn on_download_completed(&self, download_idx: usize) { + Python::with_gil(|py| { + let _ = self + .py_obj + .call_method1(py, "on_download_completed", (download_idx,)); + }); + } + + fn on_populate_cache_complete(&self, cache_entry: usize) { + Python::with_gil(|py| { + let _ = self + .py_obj + .call_method1(py, "on_populate_cache_complete", (cache_entry,)); + }); + } + + fn on_unlink_start(&self, operation: usize, record: &PrefixRecord) -> usize { + let name = record + .repodata_record + .package_record + .name + .as_normalized() + .to_string(); + Python::with_gil(|py| { + match self + .py_obj + .call_method1(py, "on_unlink_start", (operation, name)) + { + Ok(val) => val.extract::(py).unwrap_or(0), + Err(_) => 0, + } + }) + } + + fn on_unlink_complete(&self, index: usize) { + Python::with_gil(|py| { + let _ = self.py_obj.call_method1(py, "on_unlink_complete", (index,)); + }); + } + + fn on_link_start(&self, operation: usize, record: &RepoDataRecord) -> usize { + let name = record.package_record.name.as_normalized().to_string(); + Python::with_gil(|py| { + match self + .py_obj + .call_method1(py, "on_link_start", (operation, name)) + { + Ok(val) => val.extract::(py).unwrap_or(0), + Err(_) => 0, + } + }) + } + + fn on_link_complete(&self, index: usize) { + Python::with_gil(|py| { + let _ = self.py_obj.call_method1(py, "on_link_complete", (index,)); + }); + } + + fn on_transaction_operation_complete(&self, operation: usize) { + Python::with_gil(|py| { + let _ = self + .py_obj + .call_method1(py, "on_transaction_operation_complete", (operation,)); + }); + } + + fn on_transaction_complete(&self) { + Python::with_gil(|py| { + let _ = self.py_obj.call_method0(py, "on_transaction_complete"); + }); + } + + fn on_post_link_start(&self, package_name: &str, script_path: &str) -> usize { + let pkg = package_name.to_string(); + let path = script_path.to_string(); + Python::with_gil(|py| { + match self + .py_obj + .call_method1(py, "on_post_link_start", (pkg, path)) + { + Ok(val) => val.extract::(py).unwrap_or(0), + Err(_) => 0, + } + }) + } + + fn on_post_link_complete(&self, index: usize, success: bool) { + Python::with_gil(|py| { + let _ = self + .py_obj + .call_method1(py, "on_post_link_complete", (index, success)); + }); + } + + fn on_pre_unlink_start(&self, package_name: &str, script_path: &str) -> usize { + let pkg = package_name.to_string(); + let path = script_path.to_string(); + Python::with_gil(|py| { + match self + .py_obj + .call_method1(py, "on_pre_unlink_start", (pkg, path)) + { + Ok(val) => val.extract::(py).unwrap_or(0), + Err(_) => 0, + } + }) + } + + fn on_pre_unlink_complete(&self, index: usize, success: bool) { + Python::with_gil(|py| { + let _ = self + .py_obj + .call_method1(py, "on_pre_unlink_complete", (index, success)); + }); + } +} + #[pyfunction] #[allow(clippy::too_many_arguments)] -#[pyo3(signature = (records, target_prefix, execute_link_scripts=false, show_progress=false, platform=None, client=None, cache_dir=None, installed_packages=None, reinstall_packages=None, ignored_packages=None, requested_specs=None))] +#[pyo3(signature = (records, target_prefix, execute_link_scripts=false, show_progress=false, platform=None, client=None, cache_dir=None, installed_packages=None, reinstall_packages=None, ignored_packages=None, requested_specs=None, reporter=None))] pub fn py_install<'a>( py: Python<'a>, records: Vec>, @@ -32,6 +228,7 @@ pub fn py_install<'a>( reinstall_packages: Option>, ignored_packages: Option>, requested_specs: Option>, + reporter: Option, ) -> PyResult> { let dependencies = records .into_iter() @@ -70,7 +267,11 @@ pub fn py_install<'a>( future_into_py(py, async move { let mut installer = Installer::new().with_execute_link_scripts(execute_link_scripts); - if show_progress { + if let Some(py_reporter) = reporter { + installer.set_reporter(PyReporter { + py_obj: Arc::new(py_reporter), + }); + } else if show_progress { installer.set_reporter(IndicatifReporter::builder().finish()); }