From 4e620f41324f2fb02957fe1aeeb1d58ebae23489 Mon Sep 17 00:00:00 2001 From: Tim de Jager Date: Fri, 22 Dec 2023 13:56:04 +0100 Subject: [PATCH 01/10] refactor: moved pypi installation into its own module 1. Moved files 2. Clippy fixed some small issues --- src/cli/completion.rs | 2 +- src/cli/init.rs | 2 +- src/environment.rs | 533 +----------------------------------- src/install_pypi.rs | 517 ++++++++++++++++++++++++++++++++++ src/lib.rs | 1 + src/project/manifest.rs | 10 +- src/task/executable_task.rs | 5 +- tests/project_tests.rs | 4 +- 8 files changed, 539 insertions(+), 535 deletions(-) create mode 100644 src/install_pypi.rs diff --git a/src/cli/completion.rs b/src/cli/completion.rs index 768f6fc37..bc2c71087 100644 --- a/src/cli/completion.rs +++ b/src/cli/completion.rs @@ -120,7 +120,7 @@ _arguments "${_arguments_options[@]}" \ ;; "#; - let result = replace_zsh_completion(&mut script); + let result = replace_zsh_completion(script); insta::assert_snapshot!(result); } diff --git a/src/cli/init.rs b/src/cli/init.rs index b32041436..3e71d8cc1 100644 --- a/src/cli/init.rs +++ b/src/cli/init.rs @@ -236,7 +236,7 @@ mod tests { assert_eq!(read_file_content(&file_path), original_content); // Scenario 4: Path is a folder not a file, give an error. - assert!(create_or_append_file(&dir.path(), template).is_err()); + assert!(create_or_append_file(dir.path(), template).is_err()); dir.close().unwrap(); } diff --git a/src/environment.rs b/src/environment.rs index f634aaa5a..3a6cc0bbf 100644 --- a/src/environment.rs +++ b/src/environment.rs @@ -1,36 +1,15 @@ use crate::{ - consts, default_authenticated_client, install, lock_file, prefix::Prefix, progress, - progress::ProgressBarMessageFormatter, - virtual_packages::verify_current_platform_has_required_virtual_packages, Project, + consts, default_authenticated_client, install, install_pypi, lock_file, prefix::Prefix, + progress, virtual_packages::verify_current_platform_has_required_virtual_packages, Project, }; -use futures::{stream, Stream, StreamExt, TryFutureExt, TryStreamExt}; -use indexmap::IndexSet; -use indicatif::ProgressBar; -use itertools::Itertools; use miette::{Context, IntoDiagnostic, LabeledSpan}; use crate::lock_file::lock_file_satisfies_project; use rattler::install::Transaction; -use rattler_conda_types::{Platform, PrefixRecord, RepoDataRecord}; -use rattler_lock::{CondaLock, LockedDependency}; -use rip::types::Artifact; -use rip::{ - artifacts::{ - wheel::{InstallPaths, UnpackWheelOptions}, - Wheel, - }, - index::PackageDb, - python_env::{find_distributions_in_venv, uninstall_distribution, Distribution, WheelTag}, - types::{ - ArtifactHashes, ArtifactInfo, ArtifactName, Extra, NormalizedPackageName, WheelFilename, - }, -}; -use std::collections::HashSet; -use std::{io::ErrorKind, path::Path, str::FromStr, time::Duration}; -use tokio::task::JoinError; - -/// The installer name for pypi packages installed by pixi. -const PIXI_PYPI_INSTALLER: &str = env!("CARGO_PKG_NAME"); +use rattler_conda_types::{Platform, PrefixRecord}; +use rattler_lock::CondaLock; +use rip::index::PackageDb; +use std::{io::ErrorKind, path::Path}; /// Verify the location of the prefix folder is not changed so the applied prefix path is still valid. /// Errors when there is a file system error or the path does not align with the defined prefix. @@ -200,12 +179,18 @@ pub async fn update_prefix( } // Remove python packages from a previous python distribution if the python version changed. - remove_old_python_distributions(prefix, platform, &transaction)?; + install_pypi::remove_old_python_distributions(prefix, platform, &transaction)?; // Install and/or remove python packages progress::await_in_progress( "updating python packages", - update_python_distributions(package_db, prefix, lock_file, platform, &transaction), + install_pypi::update_python_distributions( + package_db, + prefix, + lock_file, + platform, + &transaction, + ), ) .await?; @@ -221,493 +206,3 @@ pub async fn update_prefix( Ok(()) } - -/// Installs and/or remove python distributions. -async fn update_python_distributions( - package_db: &PackageDb, - prefix: &Prefix, - lock_file: &CondaLock, - platform: Platform, - transaction: &Transaction, -) -> miette::Result<()> { - // Get the python info from the transaction - let Some(python_info) = transaction.python_info.as_ref() else { - return Ok(()); - }; - - // Determine where packages would have been installed - let install_paths = InstallPaths::for_venv( - ( - python_info.short_version.0 as u32, - python_info.short_version.1 as u32, - 0, - ), - platform.is_windows(), - ); - - // Determine the current python distributions in those locations - let current_python_packages = find_distributions_in_venv(prefix.root(), &install_paths) - .into_diagnostic() - .context( - "failed to locate python packages that have not been installed as conda packages", - )?; - - // Determine the python packages that are part of the lock-file - let python_packages = lock_file - .get_packages_by_platform(platform) - .filter(|p| p.is_pypi()) - .collect_vec(); - - // Determine the python packages to remove before we start installing anything new. If the - // python version changed between installations we will have to remove any previous distribution - // regardless. - let (python_distributions_to_remove, python_distributions_to_install) = - determine_python_distributions_to_remove_and_install( - prefix.root(), - current_python_packages, - python_packages, - ); - - // Start downloading the python packages that we want in the background. - let (package_stream, package_stream_pb) = - stream_python_artifacts(package_db, python_distributions_to_install.clone()); - - // Remove python packages that need to be removed - if !python_distributions_to_remove.is_empty() { - let site_package_path = install_paths.site_packages(); - - for python_distribution in python_distributions_to_remove { - uninstall_pixi_installed_distribution(prefix, site_package_path, &python_distribution)?; - } - } - - // Install the individual python packages that we want - let package_install_pb = install_python_distributions( - prefix, - install_paths, - &prefix.root().join(&python_info.path), - package_stream, - ) - .await?; - - // Clear any pending progress bar - for pb in package_install_pb - .into_iter() - .chain(package_stream_pb.into_iter()) - { - pb.finish_and_clear(); - } - - Ok(()) -} - -/// Concurrently installs python wheels as they become available. -async fn install_python_distributions( - prefix: &Prefix, - install_paths: InstallPaths, - python_executable_path: &Path, - package_stream: impl Stream, HashSet, Wheel)>> + Sized, -) -> miette::Result> { - // Determine the number of packages that we are going to install - let len = { - let (lower_bound, upper_bound) = package_stream.size_hint(); - upper_bound.unwrap_or(lower_bound) - }; - if len == 0 { - return Ok(None); - } - - // Create a progress bar to show the progress of the installation - let pb = progress::global_multi_progress().add(ProgressBar::new(len as u64)); - pb.set_style(progress::default_progress_style()); - pb.set_prefix("unpacking wheels"); - pb.enable_steady_tick(Duration::from_millis(100)); - - // Create a message formatter to show the current operation - let message_formatter = ProgressBarMessageFormatter::new(pb.clone()); - - // Concurrently unpack the wheels as they become available in the stream. - let install_pb = pb.clone(); - package_stream - .try_for_each_concurrent(Some(20), move |(hash, extras, wheel)| { - let install_paths = install_paths.clone(); - let root = prefix.root().to_path_buf(); - let message_formatter = message_formatter.clone(); - let pb = install_pb.clone(); - let python_executable_path = python_executable_path.to_owned(); - async move { - let pb_task = message_formatter.start(wheel.name().to_string()).await; - let unpack_result = tokio::task::spawn_blocking(move || { - wheel - .unpack( - &root, - &install_paths, - &python_executable_path, - &UnpackWheelOptions { - installer: Some(PIXI_PYPI_INSTALLER.into()), - extras: Some(extras), - ..Default::default() - }, - ) - .into_diagnostic() - .and_then(|unpacked_wheel| { - if let Some(hash) = hash { - std::fs::write(unpacked_wheel.dist_info.join("HASH"), hash) - .into_diagnostic() - } else { - Ok(()) - } - }) - }) - .map_err(JoinError::try_into_panic) - .await; - - pb_task.finish().await; - pb.inc(1); - - match unpack_result { - Ok(unpack_result) => unpack_result, - Err(Ok(panic)) => std::panic::resume_unwind(panic), - Err(Err(e)) => Err(miette::miette!("{e}")), - } - } - }) - .await?; - - // Update the progress bar - pb.set_style(progress::finished_progress_style()); - pb.finish(); - - Ok(Some(pb)) -} - -/// Creates a stream which downloads the specified python packages. The stream will download the -/// packages in parallel and yield them as soon as they become available. -fn stream_python_artifacts<'a>( - package_db: &'a PackageDb, - packages_to_download: Vec<&'a LockedDependency>, -) -> ( - impl Stream, HashSet, Wheel)>> + 'a, - Option, -) { - if packages_to_download.is_empty() { - return (stream::empty().left_stream(), None); - } - - // Construct a progress bar to provide some indication on what is currently downloading. - // TODO: It would be much nicer if we can provide more information with regards to the progress. - // For instance if we could also show at what speed the downloads are progressing or the total - // size of the downloads that would really help the user I think. - let pb = - progress::global_multi_progress().add(ProgressBar::new(packages_to_download.len() as u64)); - pb.set_style(progress::default_progress_style()); - pb.set_prefix("acquiring wheels"); - pb.enable_steady_tick(Duration::from_millis(100)); - - // Construct a message formatter - let message_formatter = ProgressBarMessageFormatter::new(pb.clone()); - - let stream_pb = pb.clone(); - let total_packages = packages_to_download.len(); - let download_stream = stream::iter(packages_to_download) - .map(move |package| { - let pb = stream_pb.clone(); - let message_formatter = message_formatter.clone(); - async move { - let pip_package = package - .as_pypi() - .expect("must be a pip package at this point"); - - // Determine the filename from the - let filename = pip_package - .url - .path_segments() - .and_then(|s| s.last()) - .expect("url is missing a path"); - let name = NormalizedPackageName::from_str(&package.name) - .into_diagnostic() - .with_context(|| { - format!("'{}' is not a valid python package name", &package.name) - })?; - let wheel_name = WheelFilename::from_filename(filename, &name) - .expect("failed to convert filename to wheel filename"); - - // Log out intent to install this python package. - tracing::info!("downloading python package {filename}"); - let pb_task = message_formatter.start(filename.to_string()).await; - - // Reconstruct the ArtifactInfo from the data in the lockfile. - let artifact_info = ArtifactInfo { - filename: ArtifactName::Wheel(wheel_name), - url: pip_package.url.clone(), - hashes: pip_package.hash.as_ref().map(|hash| ArtifactHashes { - sha256: hash.sha256().cloned(), - }), - requires_python: pip_package - .requires_python - .as_ref() - .map(|p| p.parse()) - .transpose() - .expect("the lock file contains an invalid 'requires_python` field"), - dist_info_metadata: Default::default(), - yanked: Default::default(), - }; - - // TODO: Maybe we should have a cache of wheels separate from the package_db. Since a - // wheel can just be identified by its hash or url. - let wheel: Wheel = package_db.get_wheel(&artifact_info, None).await?; - - // Update the progress bar - pb_task.finish().await; - pb.inc(1); - if pb.position() == total_packages as u64 { - pb.set_style(progress::finished_progress_style()); - pb.finish(); - } - - let hash = pip_package - .hash - .as_ref() - .and_then(|h| h.sha256()) - .map(|sha256| format!("sha256-{:x}", sha256)); - - Ok(( - hash, - pip_package - .extras - .iter() - .filter_map(|e| Extra::from_str(e).ok()) - .collect(), - wheel, - )) - } - }) - .buffer_unordered(20) - .right_stream(); - - (download_stream, Some(pb)) -} - -/// If there was a previous version of python installed, remove any distribution installed in that -/// environment. -fn remove_old_python_distributions( - prefix: &Prefix, - platform: Platform, - transaction: &Transaction, -) -> miette::Result<()> { - // Determine if the current distribution is the same as the desired distribution. - let Some(previous_python_installation) = transaction.current_python_info.as_ref() else { - return Ok(()); - }; - if Some(previous_python_installation.short_version) - == transaction.python_info.as_ref().map(|p| p.short_version) - { - return Ok(()); - } - - // Determine the current python distributions in its install locations - let python_version = ( - previous_python_installation.short_version.0 as u32, - previous_python_installation.short_version.1 as u32, - 0, - ); - let install_paths = InstallPaths::for_venv(python_version, platform.is_windows()); - - // Locate the packages that are installed in the previous environment - let current_python_packages = find_distributions_in_venv(prefix.root(), &install_paths) - .into_diagnostic() - .with_context(|| format!("failed to determine the python packages installed for a previous version of python ({}.{})", python_version.0, python_version.1))? - .into_iter().filter(|d| d.installer.as_deref() != Some("conda") && d.installer.is_some()).collect_vec(); - - let pb = progress::global_multi_progress() - .add(ProgressBar::new(current_python_packages.len() as u64)); - pb.set_style(progress::default_progress_style()); - pb.set_message("removing old python packages"); - pb.enable_steady_tick(Duration::from_millis(100)); - - // Remove the python packages - let site_package_path = install_paths.site_packages(); - for python_package in current_python_packages { - pb.set_message(format!( - "{} {}", - &python_package.name, &python_package.version - )); - - uninstall_pixi_installed_distribution(prefix, site_package_path, &python_package)?; - - pb.inc(1); - } - - Ok(()) -} - -/// Uninstalls a python distribution that was previously installed by pixi. -fn uninstall_pixi_installed_distribution( - prefix: &Prefix, - site_package_path: &Path, - python_package: &Distribution, -) -> miette::Result<()> { - tracing::info!( - "uninstalling python package {}-{}", - &python_package.name, - &python_package.version - ); - let relative_dist_info = python_package - .dist_info - .strip_prefix(site_package_path) - .expect("the dist-info path must be a sub-path of the site-packages path"); - - // HACK: Also remove the HASH file that pixi writes. Ignore the error if its there. We - // should probably actually add this file to the RECORD. - let _ = std::fs::remove_file(prefix.root().join(&python_package.dist_info).join("HASH")); - - uninstall_distribution(&prefix.root().join(site_package_path), relative_dist_info) - .into_diagnostic() - .with_context(|| format!("could not uninstall python package {}-{}. Manually remove the `.pixi/env` folder and try again.", &python_package.name, &python_package.version))?; - - Ok(()) -} - -/// Determine which python packages we can leave untouched and which python packages should be -/// removed. -fn determine_python_distributions_to_remove_and_install<'p>( - prefix: &Path, - mut current_python_packages: Vec, - desired_python_packages: Vec<&'p LockedDependency>, -) -> (Vec, Vec<&'p LockedDependency>) { - // Determine the artifact tags associated with the locked dependencies. - let mut desired_python_packages = extract_locked_tags(desired_python_packages); - - // Any package that is currently installed that is not part of the locked dependencies should be - // removed. So we keep it in the `current_python_packages` list. - // Any package that is in the currently installed list that is NOT found in the lockfile is - // retained in the list to mark it for removal. - current_python_packages.retain(|current_python_packages| { - if current_python_packages.installer.is_none() { - // If this package has no installer, we can't make a reliable decision on whether to - // keep it or not. So we do not uninstall it. - return false; - } - - if let Some(found_desired_packages_idx) = - desired_python_packages - .iter() - .position(|(pkg, artifact_name)| { - does_installed_match_locked_package( - prefix, - current_python_packages, - (pkg, artifact_name.as_ref()), - ) - }) - { - // Remove from the desired list of packages to install & from the packages to uninstall. - desired_python_packages.remove(found_desired_packages_idx); - false - } else { - // Only if this package was previously installed by us do we remove it. - current_python_packages.installer.as_deref() == Some(PIXI_PYPI_INSTALLER) - } - }); - - ( - current_python_packages, - desired_python_packages - .into_iter() - .map(|(pkg, _)| pkg) - .collect(), - ) -} - -/// Determine the wheel tags for the locked dependencies. These are extracted by looking at the url -/// of the locked dependency. The filename of the URL is converted to a wheel name and the tags are -/// extract from that. -/// -/// If the locked dependency is not a wheel distribution `None` is returned for the tags. If the -/// the wheel name could not be parsed `None` is returned for the tags and a warning is emitted. -fn extract_locked_tags( - desired_python_packages: Vec<&LockedDependency>, -) -> Vec<(&LockedDependency, Option>)> { - desired_python_packages - .into_iter() - .map(|pkg| { - // Get the package as a pip package. If the package is not a pip package we can just ignore it. - let Some(pip) = pkg.as_pypi() else { return (pkg, None); }; - - // Extract the filename from the url and the name from the package name. - let Some(filename) = pip.url.path_segments().and_then(|s| s.last()) else { - tracing::warn!( - "failed to determine the artifact name of the python package {}-{} from url {}: the url has no filename.", - &pkg.name, pkg.version, &pip.url); - return (pkg, None); - }; - let Ok(name) = NormalizedPackageName::from_str(&pkg.name) else { - tracing::warn!( - "failed to determine the artifact name of the python package {}-{} from url {}: {} is not a valid package name.", - &pkg.name, pkg.version, &pip.url, &pkg.name); - return (pkg, None); - }; - - // Determine the artifact type from the name and filename - match ArtifactName::from_filename(filename, &name) { - Ok(ArtifactName::Wheel(name)) => (pkg, Some(IndexSet::from_iter(name.all_tags_iter()))), - Ok(_) => (pkg, None), - Err(err) => { - tracing::warn!( - "failed to determine the artifact name of the python package {}-{}. Could not determine the name from the url {}: {err}", - &pkg.name, pkg.version, &pip.url); - (pkg, None) - } - } - }) - .collect() -} - -/// Returns true if the installed python package matches the locked python package. If that is the -/// case we can assume that the locked python package is already installed. -fn does_installed_match_locked_package( - prefix_root: &Path, - installed_python_package: &Distribution, - locked_python_package: (&LockedDependency, Option<&IndexSet>), -) -> bool { - let (pkg, artifact_tags) = locked_python_package; - - // Match on name and version - if pkg.name != installed_python_package.name.as_str() - || pep440_rs::Version::from_str(&pkg.version).ok().as_ref() - != Some(&installed_python_package.version) - { - return false; - } - - // If this distribution is installed with pixi we can assume that there is a URL file that - // contains the original URL. - if installed_python_package.installer.as_deref() == Some("pixi") { - let expected_hash = pkg - .as_pypi() - .and_then(|hash| hash.hash.as_ref()) - .and_then(|hash| hash.sha256()) - .map(|sha256| format!("sha256-{:x}", sha256)); - if let Some(expected_hash) = expected_hash { - let hash_path = prefix_root - .join(&installed_python_package.dist_info) - .join("HASH"); - if let Ok(actual_hash) = std::fs::read_to_string(hash_path) { - return actual_hash == expected_hash; - } - } - } - - // Try to match the tags of both packages. This turns out to be pretty unreliable because - // there are many WHEELS that do not report the tags of their filename correctly in the - // WHEEL file. - match (artifact_tags, &installed_python_package.tags) { - (None, _) | (_, None) => { - // One, or both, of the artifacts are not a wheel distribution so we cannot - // currently compare them. In that case we always just reinstall. - // TODO: Maybe log some info here? - // TODO: Add support for more distribution types. - false - } - (Some(locked_tags), Some(installed_tags)) => locked_tags == installed_tags, - } -} diff --git a/src/install_pypi.rs b/src/install_pypi.rs new file mode 100644 index 000000000..73d3a27c9 --- /dev/null +++ b/src/install_pypi.rs @@ -0,0 +1,517 @@ +use crate::prefix::Prefix; +use crate::progress; +use crate::progress::ProgressBarMessageFormatter; +use futures::{stream, Stream, StreamExt, TryFutureExt, TryStreamExt}; +use indexmap::IndexSet; +use indicatif::ProgressBar; +use itertools::Itertools; +use miette::{IntoDiagnostic, WrapErr}; +use rattler::install::Transaction; +use rattler_conda_types::{Platform, PrefixRecord, RepoDataRecord}; +use rattler_lock::{CondaLock, LockedDependency}; +use rip::artifacts::wheel::{InstallPaths, UnpackWheelOptions}; +use rip::artifacts::Wheel; +use rip::index::PackageDb; +use rip::python_env::{find_distributions_in_venv, uninstall_distribution, Distribution, WheelTag}; +use rip::types::{ + Artifact, ArtifactHashes, ArtifactInfo, ArtifactName, Extra, NormalizedPackageName, + WheelFilename, +}; +use std::collections::HashSet; +use std::path::Path; +use std::str::FromStr; +use std::time::Duration; +use tokio::task::JoinError; + +/// The installer name for pypi packages installed by pixi. +pub(crate) const PIXI_PYPI_INSTALLER: &str = env!("CARGO_PKG_NAME"); + +/// Installs and/or remove python distributions. +pub async fn update_python_distributions( + package_db: &PackageDb, + prefix: &Prefix, + lock_file: &CondaLock, + platform: Platform, + transaction: &Transaction, +) -> miette::Result<()> { + // Get the python info from the transaction + let Some(python_info) = transaction.python_info.as_ref() else { + return Ok(()); + }; + + // Determine where packages would have been installed + let install_paths = InstallPaths::for_venv( + ( + python_info.short_version.0 as u32, + python_info.short_version.1 as u32, + 0, + ), + platform.is_windows(), + ); + + // Determine the current python distributions in those locations + let current_python_packages = find_distributions_in_venv(prefix.root(), &install_paths) + .into_diagnostic() + .context( + "failed to locate python packages that have not been installed as conda packages", + )?; + + // Determine the python packages that are part of the lock-file + let python_packages = lock_file + .get_packages_by_platform(platform) + .filter(|p| p.is_pypi()) + .collect_vec(); + + // Determine the python packages to remove before we start installing anything new. If the + // python version changed between installations we will have to remove any previous distribution + // regardless. + let (python_distributions_to_remove, python_distributions_to_install) = + determine_python_distributions_to_remove_and_install( + prefix.root(), + current_python_packages, + python_packages, + ); + + // Start downloading the python packages that we want in the background. + let (package_stream, package_stream_pb) = + stream_python_artifacts(package_db, python_distributions_to_install.clone()); + + // Remove python packages that need to be removed + if !python_distributions_to_remove.is_empty() { + let site_package_path = install_paths.site_packages(); + + for python_distribution in python_distributions_to_remove { + uninstall_pixi_installed_distribution(prefix, site_package_path, &python_distribution)?; + } + } + + // Install the individual python packages that we want + let package_install_pb = install_python_distributions( + prefix, + install_paths, + &prefix.root().join(&python_info.path), + package_stream, + ) + .await?; + + // Clear any pending progress bar + for pb in package_install_pb + .into_iter() + .chain(package_stream_pb.into_iter()) + { + pb.finish_and_clear(); + } + + Ok(()) +} + +/// Concurrently installs python wheels as they become available. +async fn install_python_distributions( + prefix: &Prefix, + install_paths: InstallPaths, + python_executable_path: &Path, + package_stream: impl Stream, HashSet, Wheel)>> + Sized, +) -> miette::Result> { + // Determine the number of packages that we are going to install + let len = { + let (lower_bound, upper_bound) = package_stream.size_hint(); + upper_bound.unwrap_or(lower_bound) + }; + if len == 0 { + return Ok(None); + } + + // Create a progress bar to show the progress of the installation + let pb = progress::global_multi_progress().add(ProgressBar::new(len as u64)); + pb.set_style(progress::default_progress_style()); + pb.set_prefix("unpacking wheels"); + pb.enable_steady_tick(Duration::from_millis(100)); + + // Create a message formatter to show the current operation + let message_formatter = ProgressBarMessageFormatter::new(pb.clone()); + + // Concurrently unpack the wheels as they become available in the stream. + let install_pb = pb.clone(); + package_stream + .try_for_each_concurrent(Some(20), move |(hash, extras, wheel)| { + let install_paths = install_paths.clone(); + let root = prefix.root().to_path_buf(); + let message_formatter = message_formatter.clone(); + let pb = install_pb.clone(); + let python_executable_path = python_executable_path.to_owned(); + async move { + let pb_task = message_formatter.start(wheel.name().to_string()).await; + let unpack_result = tokio::task::spawn_blocking(move || { + wheel + .unpack( + &root, + &install_paths, + &python_executable_path, + &UnpackWheelOptions { + installer: Some(PIXI_PYPI_INSTALLER.into()), + extras: Some(extras), + ..Default::default() + }, + ) + .into_diagnostic() + .and_then(|unpacked_wheel| { + if let Some(hash) = hash { + std::fs::write(unpacked_wheel.dist_info.join("HASH"), hash) + .into_diagnostic() + } else { + Ok(()) + } + }) + }) + .map_err(JoinError::try_into_panic) + .await; + + pb_task.finish().await; + pb.inc(1); + + match unpack_result { + Ok(unpack_result) => unpack_result, + Err(Ok(panic)) => std::panic::resume_unwind(panic), + Err(Err(e)) => Err(miette::miette!("{e}")), + } + } + }) + .await?; + + // Update the progress bar + pb.set_style(progress::finished_progress_style()); + pb.finish(); + + Ok(Some(pb)) +} + +/// Creates a stream which downloads the specified python packages. The stream will download the +/// packages in parallel and yield them as soon as they become available. +fn stream_python_artifacts<'a>( + package_db: &'a PackageDb, + packages_to_download: Vec<&'a LockedDependency>, +) -> ( + impl Stream, HashSet, Wheel)>> + 'a, + Option, +) { + if packages_to_download.is_empty() { + return (stream::empty().left_stream(), None); + } + + // Construct a progress bar to provide some indication on what is currently downloading. + // TODO: It would be much nicer if we can provide more information with regards to the progress. + // For instance if we could also show at what speed the downloads are progressing or the total + // size of the downloads that would really help the user I think. + let pb = + progress::global_multi_progress().add(ProgressBar::new(packages_to_download.len() as u64)); + pb.set_style(progress::default_progress_style()); + pb.set_prefix("acquiring wheels"); + pb.enable_steady_tick(Duration::from_millis(100)); + + // Construct a message formatter + let message_formatter = ProgressBarMessageFormatter::new(pb.clone()); + + let stream_pb = pb.clone(); + let total_packages = packages_to_download.len(); + let download_stream = stream::iter(packages_to_download) + .map(move |package| { + let pb = stream_pb.clone(); + let message_formatter = message_formatter.clone(); + async move { + let pip_package = package + .as_pypi() + .expect("must be a pip package at this point"); + + // Determine the filename from the + let filename = pip_package + .url + .path_segments() + .and_then(|s| s.last()) + .expect("url is missing a path"); + let name = NormalizedPackageName::from_str(&package.name) + .into_diagnostic() + .with_context(|| { + format!("'{}' is not a valid python package name", &package.name) + })?; + let wheel_name = WheelFilename::from_filename(filename, &name) + .expect("failed to convert filename to wheel filename"); + + // Log out intent to install this python package. + tracing::info!("downloading python package {filename}"); + let pb_task = message_formatter.start(filename.to_string()).await; + + // Reconstruct the ArtifactInfo from the data in the lockfile. + let artifact_info = ArtifactInfo { + filename: ArtifactName::Wheel(wheel_name), + url: pip_package.url.clone(), + hashes: pip_package.hash.as_ref().map(|hash| ArtifactHashes { + sha256: hash.sha256().cloned(), + }), + requires_python: pip_package + .requires_python + .as_ref() + .map(|p| p.parse()) + .transpose() + .expect("the lock file contains an invalid 'requires_python` field"), + dist_info_metadata: Default::default(), + yanked: Default::default(), + }; + + // TODO: Maybe we should have a cache of wheels separate from the package_db. Since a + // wheel can just be identified by its hash or url. + let wheel: Wheel = package_db.get_wheel(&artifact_info, None).await?; + + // Update the progress bar + pb_task.finish().await; + pb.inc(1); + if pb.position() == total_packages as u64 { + pb.set_style(progress::finished_progress_style()); + pb.finish(); + } + + let hash = pip_package + .hash + .as_ref() + .and_then(|h| h.sha256()) + .map(|sha256| format!("sha256-{:x}", sha256)); + + Ok(( + hash, + pip_package + .extras + .iter() + .filter_map(|e| Extra::from_str(e).ok()) + .collect(), + wheel, + )) + } + }) + .buffer_unordered(20) + .right_stream(); + + (download_stream, Some(pb)) +} + +/// If there was a previous version of python installed, remove any distribution installed in that +/// environment. +pub fn remove_old_python_distributions( + prefix: &Prefix, + platform: Platform, + transaction: &Transaction, +) -> miette::Result<()> { + // Determine if the current distribution is the same as the desired distribution. + let Some(previous_python_installation) = transaction.current_python_info.as_ref() else { + return Ok(()); + }; + if Some(previous_python_installation.short_version) + == transaction.python_info.as_ref().map(|p| p.short_version) + { + return Ok(()); + } + + // Determine the current python distributions in its install locations + let python_version = ( + previous_python_installation.short_version.0 as u32, + previous_python_installation.short_version.1 as u32, + 0, + ); + let install_paths = InstallPaths::for_venv(python_version, platform.is_windows()); + + // Locate the packages that are installed in the previous environment + let current_python_packages = find_distributions_in_venv(prefix.root(), &install_paths) + .into_diagnostic() + .with_context(|| format!("failed to determine the python packages installed for a previous version of python ({}.{})", python_version.0, python_version.1))? + .into_iter().filter(|d| d.installer.as_deref() != Some("conda") && d.installer.is_some()).collect_vec(); + + let pb = progress::global_multi_progress() + .add(ProgressBar::new(current_python_packages.len() as u64)); + pb.set_style(progress::default_progress_style()); + pb.set_message("removing old python packages"); + pb.enable_steady_tick(Duration::from_millis(100)); + + // Remove the python packages + let site_package_path = install_paths.site_packages(); + for python_package in current_python_packages { + pb.set_message(format!( + "{} {}", + &python_package.name, &python_package.version + )); + + uninstall_pixi_installed_distribution(prefix, site_package_path, &python_package)?; + + pb.inc(1); + } + + Ok(()) +} + +/// Uninstalls a python distribution that was previously installed by pixi. +fn uninstall_pixi_installed_distribution( + prefix: &Prefix, + site_package_path: &Path, + python_package: &Distribution, +) -> miette::Result<()> { + tracing::info!( + "uninstalling python package {}-{}", + &python_package.name, + &python_package.version + ); + let relative_dist_info = python_package + .dist_info + .strip_prefix(site_package_path) + .expect("the dist-info path must be a sub-path of the site-packages path"); + + // HACK: Also remove the HASH file that pixi writes. Ignore the error if its there. We + // should probably actually add this file to the RECORD. + let _ = std::fs::remove_file(prefix.root().join(&python_package.dist_info).join("HASH")); + + uninstall_distribution(&prefix.root().join(site_package_path), relative_dist_info) + .into_diagnostic() + .with_context(|| format!("could not uninstall python package {}-{}. Manually remove the `.pixi/env` folder and try again.", &python_package.name, &python_package.version))?; + + Ok(()) +} + +/// Determine which python packages we can leave untouched and which python packages should be +/// removed. +fn determine_python_distributions_to_remove_and_install<'p>( + prefix: &Path, + mut current_python_packages: Vec, + desired_python_packages: Vec<&'p LockedDependency>, +) -> (Vec, Vec<&'p LockedDependency>) { + // Determine the artifact tags associated with the locked dependencies. + let mut desired_python_packages = extract_locked_tags(desired_python_packages); + + // Any package that is currently installed that is not part of the locked dependencies should be + // removed. So we keep it in the `current_python_packages` list. + // Any package that is in the currently installed list that is NOT found in the lockfile is + // retained in the list to mark it for removal. + current_python_packages.retain(|current_python_packages| { + if current_python_packages.installer.is_none() { + // If this package has no installer, we can't make a reliable decision on whether to + // keep it or not. So we do not uninstall it. + return false; + } + + if let Some(found_desired_packages_idx) = + desired_python_packages + .iter() + .position(|(pkg, artifact_name)| { + does_installed_match_locked_package( + prefix, + current_python_packages, + (pkg, artifact_name.as_ref()), + ) + }) + { + // Remove from the desired list of packages to install & from the packages to uninstall. + desired_python_packages.remove(found_desired_packages_idx); + false + } else { + // Only if this package was previously installed by us do we remove it. + current_python_packages.installer.as_deref() == Some(PIXI_PYPI_INSTALLER) + } + }); + + ( + current_python_packages, + desired_python_packages + .into_iter() + .map(|(pkg, _)| pkg) + .collect(), + ) +} + +/// Determine the wheel tags for the locked dependencies. These are extracted by looking at the url +/// of the locked dependency. The filename of the URL is converted to a wheel name and the tags are +/// extract from that. +/// +/// If the locked dependency is not a wheel distribution `None` is returned for the tags. If the +/// the wheel name could not be parsed `None` is returned for the tags and a warning is emitted. +fn extract_locked_tags( + desired_python_packages: Vec<&LockedDependency>, +) -> Vec<(&LockedDependency, Option>)> { + desired_python_packages + .into_iter() + .map(|pkg| { + // Get the package as a pip package. If the package is not a pip package we can just ignore it. + let Some(pip) = pkg.as_pypi() else { return (pkg, None); }; + + // Extract the filename from the url and the name from the package name. + let Some(filename) = pip.url.path_segments().and_then(|s| s.last()) else { + tracing::warn!( + "failed to determine the artifact name of the python package {}-{} from url {}: the url has no filename.", + &pkg.name, pkg.version, &pip.url); + return (pkg, None); + }; + let Ok(name) = NormalizedPackageName::from_str(&pkg.name) else { + tracing::warn!( + "failed to determine the artifact name of the python package {}-{} from url {}: {} is not a valid package name.", + &pkg.name, pkg.version, &pip.url, &pkg.name); + return (pkg, None); + }; + + // Determine the artifact type from the name and filename + match ArtifactName::from_filename(filename, &name) { + Ok(ArtifactName::Wheel(name)) => (pkg, Some(IndexSet::from_iter(name.all_tags_iter()))), + Ok(_) => (pkg, None), + Err(err) => { + tracing::warn!( + "failed to determine the artifact name of the python package {}-{}. Could not determine the name from the url {}: {err}", + &pkg.name, pkg.version, &pip.url); + (pkg, None) + } + } + }) + .collect() +} + +/// Returns true if the installed python package matches the locked python package. If that is the +/// case we can assume that the locked python package is already installed. +fn does_installed_match_locked_package( + prefix_root: &Path, + installed_python_package: &Distribution, + locked_python_package: (&LockedDependency, Option<&IndexSet>), +) -> bool { + let (pkg, artifact_tags) = locked_python_package; + + // Match on name and version + if pkg.name != installed_python_package.name.as_str() + || pep440_rs::Version::from_str(&pkg.version).ok().as_ref() + != Some(&installed_python_package.version) + { + return false; + } + + // If this distribution is installed with pixi we can assume that there is a URL file that + // contains the original URL. + if installed_python_package.installer.as_deref() == Some("pixi") { + let expected_hash = pkg + .as_pypi() + .and_then(|hash| hash.hash.as_ref()) + .and_then(|hash| hash.sha256()) + .map(|sha256| format!("sha256-{:x}", sha256)); + if let Some(expected_hash) = expected_hash { + let hash_path = prefix_root + .join(&installed_python_package.dist_info) + .join("HASH"); + if let Ok(actual_hash) = std::fs::read_to_string(hash_path) { + return actual_hash == expected_hash; + } + } + } + + // Try to match the tags of both packages. This turns out to be pretty unreliable because + // there are many WHEELS that do not report the tags of their filename correctly in the + // WHEEL file. + match (artifact_tags, &installed_python_package.tags) { + (None, _) | (_, None) => { + // One, or both, of the artifacts are not a wheel distribution so we cannot + // currently compare them. In that case we always just reinstall. + // TODO: Maybe log some info here? + // TODO: Add support for more distribution types. + false + } + (Some(locked_tags), Some(installed_tags)) => locked_tags == installed_tags, + } +} diff --git a/src/lib.rs b/src/lib.rs index de886aa46..1df0b3c22 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -3,6 +3,7 @@ pub mod config; pub mod consts; pub mod environment; pub mod install; +pub mod install_pypi; pub mod lock_file; pub mod prefix; pub mod progress; diff --git a/src/project/manifest.rs b/src/project/manifest.rs index 504bf3a69..848fb6567 100644 --- a/src/project/manifest.rs +++ b/src/project/manifest.rs @@ -1531,9 +1531,7 @@ mod test { vec![Platform::Linux64, Platform::Win64] ); - manifest - .add_platforms(vec![Platform::OsxArm64].iter()) - .unwrap(); + manifest.add_platforms([Platform::OsxArm64].iter()).unwrap(); assert_eq!( manifest.parsed.project.platforms.value, @@ -1588,7 +1586,7 @@ mod test { assert_eq!(manifest.parsed.project.channels, vec![]); - manifest.add_channels(vec!["conda-forge"].iter()).unwrap(); + manifest.add_channels(["conda-forge"].iter()).unwrap(); assert_eq!( manifest.parsed.project.channels, @@ -1617,9 +1615,7 @@ mod test { vec![Channel::from_str("conda-forge", &ChannelConfig::default()).unwrap()] ); - manifest - .remove_channels(vec!["conda-forge"].iter()) - .unwrap(); + manifest.remove_channels(["conda-forge"].iter()).unwrap(); assert_eq!(manifest.parsed.project.channels, vec![]); } diff --git a/src/task/executable_task.rs b/src/task/executable_task.rs index 63b8e49db..6b7b6666d 100644 --- a/src/task/executable_task.rs +++ b/src/task/executable_task.rs @@ -376,9 +376,6 @@ mod tests { let task = executable_tasks.get(0).unwrap(); assert!(task.task().is_custom()); - assert_eq!( - task.task().as_single_command().unwrap(), - r###""echo bla""### - ); + assert_eq!(task.task().as_single_command().unwrap(), r#""echo bla""#); } } diff --git a/tests/project_tests.rs b/tests/project_tests.rs index 8e6b1b11b..f1edad5d3 100644 --- a/tests/project_tests.rs +++ b/tests/project_tests.rs @@ -39,9 +39,7 @@ async fn add_channel() { // Our channel should be in the list of channels let local_channel = Channel::from_str( - Url::from_directory_path(additional_channel_dir.path()) - .unwrap() - .to_string(), + &Url::from_directory_path(additional_channel_dir.path()).unwrap(), &ChannelConfig::default(), ) .unwrap(); From 346c5dde1c35097ace4c7cc9944efcd9786efac7 Mon Sep 17 00:00:00 2001 From: Tim de Jager Date: Fri, 22 Dec 2023 14:19:08 +0100 Subject: [PATCH 02/10] feat: remove mut --- src/cli/completion.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/cli/completion.rs b/src/cli/completion.rs index bc2c71087..96ad0f63d 100644 --- a/src/cli/completion.rs +++ b/src/cli/completion.rs @@ -84,7 +84,7 @@ mod tests { #[test] pub fn test_zsh_completion() { - let mut script = r#" + let script = r#" (add) _arguments "${_arguments_options[@]}" \ '--manifest-path=[The path to '\''pixi.toml'\'']:MANIFEST_PATH:_files' \ From 825cf466bf8b8b36189f1059c8e27ee02f99984d Mon Sep 17 00:00:00 2001 From: Tim de Jager Date: Mon, 25 Dec 2023 06:55:42 +0100 Subject: [PATCH 03/10] wip: split up conda and pypi lock file update --- src/cli/add.rs | 13 +- src/cli/project/channel/add.rs | 4 +- src/cli/project/channel/remove.rs | 4 +- src/cli/project/platform/add.rs | 4 +- src/cli/project/platform/remove.rs | 4 +- src/environment.rs | 5 +- src/lock_file/mod.rs | 253 ++++++++++++++++++++--------- src/lock_file/pypi.rs | 19 ++- tests/project_tests.rs | 2 +- 9 files changed, 210 insertions(+), 98 deletions(-) diff --git a/src/cli/add.rs b/src/cli/add.rs index 1720c45f9..a9fb9bc78 100644 --- a/src/cli/add.rs +++ b/src/cli/add.rs @@ -1,9 +1,10 @@ use crate::environment::{update_prefix, verify_prefix_location_unchanged}; +use crate::lock_file::update_lock_file_for_pypi; use crate::prefix::Prefix; use crate::project::{DependencyType, SpecType}; use crate::{ consts, - lock_file::{load_lock_file, update_lock_file}, + lock_file::{load_lock_file, update_lock_file_conda}, project::python::PyPiRequirement, project::Project, virtual_packages::get_minimal_virtual_packages, @@ -345,7 +346,15 @@ async fn update_environment( ) -> miette::Result<()> { // Update the lock file let lock_file = if !no_update_lockfile { - Some(update_lock_file(project, load_lock_file(project).await?, sparse_repo_data).await?) + let lock = + update_lock_file_conda(project, load_lock_file(project).await?, sparse_repo_data) + .await?; + + if project.has_pypi_dependencies() { + update_lock_file_for_pypi(project, lock).await?.into() + } else { + lock.into() + } } else { None }; diff --git a/src/cli/project/channel/add.rs b/src/cli/project/channel/add.rs index b42f61804..4993d6d0a 100644 --- a/src/cli/project/channel/add.rs +++ b/src/cli/project/channel/add.rs @@ -1,5 +1,5 @@ use crate::environment::update_prefix; -use crate::lock_file::{load_lock_file, update_lock_file}; +use crate::lock_file::{load_lock_file, update_lock_file_conda}; use crate::prefix::Prefix; use crate::Project; use clap::Parser; @@ -52,7 +52,7 @@ pub async fn execute(mut project: Project, args: Args) -> miette::Result<()> { .add_channels(missing_channels.iter().map(|(name, _channel)| name))?; // Try to update the lock-file with the new channels - let lock_file = update_lock_file(&project, lock_file, None).await?; + let lock_file = update_lock_file_conda(&project, lock_file, None).await?; project.save()?; // Update the installation if needed diff --git a/src/cli/project/channel/remove.rs b/src/cli/project/channel/remove.rs index 8fa432b22..90b9c60e6 100644 --- a/src/cli/project/channel/remove.rs +++ b/src/cli/project/channel/remove.rs @@ -1,5 +1,5 @@ use crate::environment::update_prefix; -use crate::lock_file::{load_lock_file, update_lock_file}; +use crate::lock_file::{load_lock_file, update_lock_file_conda}; use crate::prefix::Prefix; use crate::Project; use clap::Parser; @@ -52,7 +52,7 @@ pub async fn execute(mut project: Project, args: Args) -> miette::Result<()> { .remove_channels(channels_to_remove.iter().map(|(name, _channel)| name))?; // Try to update the lock-file without the removed channels - let lock_file = update_lock_file(&project, lock_file, None).await?; + let lock_file = update_lock_file_conda(&project, lock_file, None).await?; project.save()?; // Update the installation if needed diff --git a/src/cli/project/platform/add.rs b/src/cli/project/platform/add.rs index 4d57e07c6..db0df24fc 100644 --- a/src/cli/project/platform/add.rs +++ b/src/cli/project/platform/add.rs @@ -1,7 +1,7 @@ use std::str::FromStr; use crate::environment::update_prefix; -use crate::lock_file::{load_lock_file, update_lock_file}; +use crate::lock_file::{load_lock_file, update_lock_file_conda}; use crate::prefix::Prefix; use crate::Project; use clap::Parser; @@ -49,7 +49,7 @@ pub async fn execute(mut project: Project, args: Args) -> miette::Result<()> { project.manifest.add_platforms(missing_platforms.iter())?; // Try to update the lock-file with the new channels - let lock_file = update_lock_file(&project, lock_file, None).await?; + let lock_file = update_lock_file_conda(&project, lock_file, None).await?; project.save()?; // Update the installation if needed diff --git a/src/cli/project/platform/remove.rs b/src/cli/project/platform/remove.rs index b754a1a47..5fb7de9af 100644 --- a/src/cli/project/platform/remove.rs +++ b/src/cli/project/platform/remove.rs @@ -1,5 +1,5 @@ use crate::environment::update_prefix; -use crate::lock_file::{load_lock_file, update_lock_file}; +use crate::lock_file::{load_lock_file, update_lock_file_conda}; use crate::prefix::Prefix; use crate::Project; use clap::Parser; @@ -50,7 +50,7 @@ pub async fn execute(mut project: Project, args: Args) -> miette::Result<()> { .remove_platforms(platforms_to_remove.iter().map(|p| p.to_string()))?; // Try to update the lock-file without the removed platform(s) - let lock_file = update_lock_file(&project, lock_file, None).await?; + let lock_file = update_lock_file_conda(&project, lock_file, None).await?; project.save()?; // Update the installation if needed diff --git a/src/environment.rs b/src/environment.rs index 3a6cc0bbf..1bc2832db 100644 --- a/src/environment.rs +++ b/src/environment.rs @@ -121,7 +121,10 @@ pub async fn get_up_to_date_prefix( match usage { LockFileUsage::Update => { if !up_to_date { - lock_file = lock_file::update_lock_file(project, lock_file, None).await? + lock_file = lock_file::update_lock_file_conda(project, lock_file, None).await?; + if project.has_pypi_dependencies() { + lock_file = lock_file::update_lock_file_for_pypi(project, lock_file).await?; + } } } LockFileUsage::Locked => { diff --git a/src/lock_file/mod.rs b/src/lock_file/mod.rs index bc3bba5db..fc983fd02 100644 --- a/src/lock_file/mod.rs +++ b/src/lock_file/mod.rs @@ -7,7 +7,7 @@ use crate::{progress, Project}; use futures::TryStreamExt; use futures::{stream, StreamExt}; use indicatif::ProgressBar; -use itertools::Itertools; +use itertools::{izip, Itertools}; use miette::{Context, IntoDiagnostic}; use rattler_conda_types::{ GenericVirtualPackage, MatchSpec, PackageName, Platform, RepoDataRecord, @@ -39,8 +39,36 @@ pub async fn load_lock_file(project: &Project) -> miette::Result { .unwrap_or_else(|e| Err(e).into_diagnostic()) } -/// Updates the lock file for a project. -pub async fn update_lock_file( +fn main_progress_bar(num_bars: u64, message: &'static str) -> ProgressBar { + let multi_progress = progress::global_multi_progress(); + let top_level_progress = multi_progress.add(ProgressBar::new(num_bars)); + top_level_progress.set_style(progress::long_running_progress_style()); + top_level_progress.set_message(message); + top_level_progress.enable_steady_tick(Duration::from_millis(50)); + top_level_progress +} + +fn platform_solve_bars(platforms: &[Platform]) -> Vec { + platforms + .iter() + .map(|platform| { + let pb = + progress::global_multi_progress().add(ProgressBar::new(platforms.len() as u64)); + pb.set_style( + indicatif::ProgressStyle::with_template(&format!( + " {:<9} ..", + platform.to_string(), + )) + .unwrap(), + ); + pb.enable_steady_tick(Duration::from_millis(100)); + pb + }) + .collect_vec() +} + +/// Updates the lock file for conda dependencies for the specified project. +pub async fn update_lock_file_conda( project: &Project, existing_lock_file: CondaLock, repodata: Option>, @@ -56,11 +84,10 @@ pub async fn update_lock_file( .into(); // Construct a progress bar - let multi_progress = progress::global_multi_progress(); - let top_level_progress = multi_progress.add(ProgressBar::new(platforms.len() as u64)); - top_level_progress.set_style(progress::long_running_progress_style()); - top_level_progress.set_message("solving dependencies"); - top_level_progress.enable_steady_tick(Duration::from_millis(50)); + let _top_level_progress = + main_progress_bar(platforms.len() as u64, "resolving conda dependencies"); + // Create progress bars for each platform + let solve_bars = platform_solve_bars(platforms); // Construct a conda lock file let channels = project @@ -68,35 +95,6 @@ pub async fn update_lock_file( .iter() .map(|channel| rattler_lock::Channel::from(channel.base_url().to_string())); - // Create progress bars for each platform - let solve_bars = platforms - .iter() - .map(|platform| { - let pb = - progress::global_multi_progress().add(ProgressBar::new(platforms.len() as u64)); - pb.set_style( - indicatif::ProgressStyle::with_template(&format!( - " {:<9} ..", - platform.to_string(), - )) - .unwrap(), - ); - pb.enable_steady_tick(Duration::from_millis(100)); - pb - }) - .collect_vec(); - - // Solve each platform concurrently - let num_concurrent = if project.has_pypi_dependencies() { - // HACK: There is a bug in rip that causes a dead-lock when solving multiple environments - // at the same time. So if there are pypi dependencies we limit the number of concurrent - // solves to 1. - 1 - } else { - // By default we solve 2 platforms concurrently. Could probably do more but solving takes - // a significant amount of memory. - 2 - }; let result: miette::Result> = stream::iter(platforms.iter().zip(solve_bars.iter().cloned())) .map(|(platform, pb)| { @@ -134,7 +132,7 @@ pub async fn update_lock_file( Ok(result) } }) - .buffer_unordered(num_concurrent) + .buffer_unordered(2) .try_collect() .await; @@ -157,6 +155,139 @@ pub async fn update_lock_file( Ok(conda_lock) } +pub async fn update_lock_file_for_pypi( + project: &Project, + new_lock_file: CondaLock, +) -> miette::Result { + let platforms = project.platforms(); + let _top_level_progress = + main_progress_bar(platforms.len() as u64, "resolving pypi dependencies"); + let solve_bars = platform_solve_bars(platforms); + + let records = platforms + .iter() + .map(|plat| new_lock_file.get_conda_packages_by_platform(*plat)); + + let result: miette::Result> = + stream::iter(izip!(platforms.iter(), solve_bars.iter().cloned(), records)) + .map(|(platform, pb, records)| { + pb.reset_elapsed(); + pb.set_style( + indicatif::ProgressStyle::with_template(&format!( + " {{spinner:.dim}} {:<9} [{{elapsed_precise}}] {{msg:.dim}}", + platform.to_string(), + )) + .unwrap(), + ); + + async move { + let locked_packages = LockedPackagesBuilder::new(*platform); + let result = resolve_pypi( + project, + &records.into_diagnostic()?, + locked_packages, + *platform, + &pb, + ) + .await?; + + pb.set_style( + indicatif::ProgressStyle::with_template(&format!( + " {} {:<9} [{{elapsed_precise}}]", + console::style(console::Emoji("✔", "↳")).green(), + platform.to_string(), + )) + .unwrap(), + ); + pb.finish(); + + Ok(result) + } + }) + // TODO: Hack to ensure we do not encounter file-locking issues in windows, should look at a better solution + .buffer_unordered(1) + .try_collect() + .await; + + // Clear all progress bars + for bar in solve_bars { + bar.finish_and_clear(); + } + + let channels = project + .channels() + .iter() + .map(|channel| rattler_lock::Channel::from(channel.base_url().to_string())); + let mut builder = LockFileBuilder::new(channels, platforms.iter().cloned(), vec![]); + for locked_packages in result? { + builder = builder.add_locked_packages(locked_packages); + } + let conda_lock = builder.build().into_diagnostic()?; + + // TODO: think of a better way to do this + Ok(CondaLock { + metadata: new_lock_file.metadata, + package: Vec::from_iter( + conda_lock + .package + .iter() + .chain(new_lock_file.package.iter()) + .cloned(), + ), + }) +} + +async fn resolve_pypi( + project: &Project, + records: &[RepoDataRecord], + mut locked_packages: LockedPackagesBuilder, + platform: Platform, + pb: &ProgressBar, +) -> miette::Result { + // Solve python packages + pb.set_message("resolving python"); + let python_artifacts = pypi::resolve_dependencies(project, platform, records).await?; + + // Clear message + pb.set_message(""); + + // Add pip packages + for python_artifact in python_artifacts { + let (artifact, metadata) = project + .pypi_package_db()? + .get_metadata(&python_artifact.artifacts, None) + .await + .expect("failed to get metadata for a package for which we have already fetched metadata during solving.") + .expect("no metadata for a package for which we have already fetched metadata during solving."); + + let locked_package = PypiLockedDependencyBuilder { + name: python_artifact.name.to_string(), + version: python_artifact.version.to_string(), + requires_dist: metadata + .requires_dist + .into_iter() + .map(|r| r.to_string()) + .collect(), + requires_python: metadata.requires_python.map(|r| r.to_string()), + extras: python_artifact + .extras + .into_iter() + .map(|e| e.as_str().to_string()) + .collect(), + url: artifact.url.clone(), + hash: artifact + .hashes + .as_ref() + .and_then(|hash| PackageHashes::from_hashes(None, hash.sha256)), + source: None, + build: None, + }; + + locked_packages.add_locked_package(locked_package) + } + Ok(locked_packages) +} + async fn resolve_platform( project: &Project, existing_lock_file: &CondaLock, @@ -198,56 +329,22 @@ async fn resolve_platform( ) .await?; - // Solve python packages - pb.set_message("resolving python"); - let python_artifacts = pypi::resolve_pypi_dependencies(project, platform, &mut records).await?; - - // Clear message - pb.set_message(""); + // Add purl's for the conda packages that are also available as pypi packages + pypi::amend_pypi_purls(&mut records).await?; // Update lock file let mut locked_packages = LockedPackagesBuilder::new(platform); // Add conda packages - for record in records { + for record in records.iter() { let locked_package = CondaLockedDependencyBuilder::try_from(record).into_diagnostic()?; locked_packages.add_locked_package(locked_package); } - // Add pip packages - for python_artifact in python_artifacts { - let (artifact, metadata) = project - .pypi_package_db()? - .get_metadata(&python_artifact.artifacts, None) - .await - .expect("failed to get metadata for a package for which we have already fetched metadata during solving.") - .expect("no metadata for a package for which we have already fetched metadata during solving."); + // Add pypi packages + // let locked_packages = + // resolve_pypi(project, &mut records, locked_packages, platform, &pb).await?; - let locked_package = PypiLockedDependencyBuilder { - name: python_artifact.name.to_string(), - version: python_artifact.version.to_string(), - requires_dist: metadata - .requires_dist - .into_iter() - .map(|r| r.to_string()) - .collect(), - requires_python: metadata.requires_python.map(|r| r.to_string()), - extras: python_artifact - .extras - .into_iter() - .map(|e| e.as_str().to_string()) - .collect(), - url: artifact.url.clone(), - hash: artifact - .hashes - .as_ref() - .and_then(|hash| PackageHashes::from_hashes(None, hash.sha256)), - source: None, - build: None, - }; - - locked_packages.add_locked_package(locked_package) - } Ok(locked_packages) } diff --git a/src/lock_file/pypi.rs b/src/lock_file/pypi.rs index 8cd6c6640..ccc4318de 100644 --- a/src/lock_file/pypi.rs +++ b/src/lock_file/pypi.rs @@ -14,22 +14,16 @@ use rip::resolve::{resolve, PinnedPackage, ResolveOptions, SDistResolution}; use std::{collections::HashMap, str::FromStr, vec}; /// Resolve python packages for the specified project. -pub async fn resolve_pypi_dependencies<'p>( +pub async fn resolve_dependencies<'p>( project: &'p Project, platform: Platform, - conda_packages: &mut [RepoDataRecord], + conda_packages: &[RepoDataRecord], ) -> miette::Result>> { let dependencies = project.pypi_dependencies(platform); if dependencies.is_empty() { return Ok(vec![]); } - // Amend the records with pypi purls if they are not present yet. - let conda_forge_mapping = pypi_name_mapping::conda_pypi_name_mapping().await?; - for record in conda_packages.iter_mut() { - pypi_name_mapping::amend_pypi_purls(record, conda_forge_mapping)?; - } - // Determine the python packages that are installed by the conda packages let conda_python_packages = package_identifier::PypiPackageIdentifier::from_records(conda_packages) @@ -100,6 +94,15 @@ pub async fn resolve_pypi_dependencies<'p>( Ok(result) } +/// Amend the records with pypi purls if they are not present yet. +pub async fn amend_pypi_purls(conda_packages: &mut [RepoDataRecord]) -> miette::Result<()> { + let conda_forge_mapping = pypi_name_mapping::conda_pypi_name_mapping().await?; + for record in conda_packages.iter_mut() { + pypi_name_mapping::amend_pypi_purls(record, conda_forge_mapping)?; + } + Ok(()) +} + /// Returns true if the specified record refers to a version/variant of python. pub fn is_python_record(record: &RepoDataRecord) -> bool { package_name_is_python(&record.package_record.name) diff --git a/tests/project_tests.rs b/tests/project_tests.rs index f1edad5d3..aae914442 100644 --- a/tests/project_tests.rs +++ b/tests/project_tests.rs @@ -39,7 +39,7 @@ async fn add_channel() { // Our channel should be in the list of channels let local_channel = Channel::from_str( - &Url::from_directory_path(additional_channel_dir.path()).unwrap(), + Url::from_directory_path(additional_channel_dir.path()).unwrap(), &ChannelConfig::default(), ) .unwrap(); From e7b86269db4a8965c9e783f334c03d5bf921e777 Mon Sep 17 00:00:00 2001 From: Tim de Jager Date: Mon, 1 Jan 2024 20:58:14 +0100 Subject: [PATCH 04/10] feat: split pypi and conda environment creation --- src/cli/add.rs | 12 +-- src/cli/install.rs | 2 +- src/cli/project/channel/add.rs | 31 ++----- src/cli/project/channel/remove.rs | 29 +------ src/cli/project/platform/add.rs | 26 +----- src/cli/project/platform/remove.rs | 28 +----- src/cli/remove.rs | 2 +- src/cli/run.rs | 2 +- src/cli/shell.rs | 2 +- src/environment.rs | 131 +++++++++++++++++++---------- src/install_pypi.rs | 49 +++++------ src/lock_file/mod.rs | 19 +++-- 12 files changed, 144 insertions(+), 189 deletions(-) diff --git a/src/cli/add.rs b/src/cli/add.rs index a9fb9bc78..21939e9f4 100644 --- a/src/cli/add.rs +++ b/src/cli/add.rs @@ -1,4 +1,4 @@ -use crate::environment::{update_prefix, verify_prefix_location_unchanged}; +use crate::environment::{update_prefix_conda, verify_prefix_location_unchanged}; use crate::lock_file::update_lock_file_for_pypi; use crate::prefix::Prefix; use crate::project::{DependencyType, SpecType}; @@ -368,14 +368,8 @@ async fn update_environment( let installed_packages = prefix.find_installed_packages(None).await?; // Update the prefix - update_prefix( - project.pypi_package_db()?, - &prefix, - installed_packages, - &lock_file, - Platform::current(), - ) - .await?; + update_prefix_conda(&prefix, installed_packages, &lock_file, Platform::current()) + .await?; } } Ok(()) diff --git a/src/cli/install.rs b/src/cli/install.rs index a7225035e..ac475b911 100644 --- a/src/cli/install.rs +++ b/src/cli/install.rs @@ -17,7 +17,7 @@ pub struct Args { pub async fn execute(args: Args) -> miette::Result<()> { let project = Project::load_or_else_discover(args.manifest_path.as_deref())?; - get_up_to_date_prefix(&project, args.lock_file_usage.into()).await?; + get_up_to_date_prefix(&project, args.lock_file_usage.into(), false).await?; // Emit success eprintln!( diff --git a/src/cli/project/channel/add.rs b/src/cli/project/channel/add.rs index 4993d6d0a..40aed24f6 100644 --- a/src/cli/project/channel/add.rs +++ b/src/cli/project/channel/add.rs @@ -1,11 +1,11 @@ -use crate::environment::update_prefix; -use crate::lock_file::{load_lock_file, update_lock_file_conda}; -use crate::prefix::Prefix; +use crate::environment::{get_up_to_date_prefix, LockFileUsage}; +use crate::lock_file::load_lock_file; + use crate::Project; use clap::Parser; use itertools::Itertools; use miette::IntoDiagnostic; -use rattler_conda_types::{Channel, ChannelConfig, Platform}; +use rattler_conda_types::{Channel, ChannelConfig}; #[derive(Parser, Debug, Default)] pub struct Args { @@ -44,34 +44,15 @@ pub async fn execute(mut project: Project, args: Args) -> miette::Result<()> { } // Load the existing lock-file - let lock_file = load_lock_file(&project).await?; + let _lock_file = load_lock_file(&project).await?; // Add the channels to the lock-file project .manifest .add_channels(missing_channels.iter().map(|(name, _channel)| name))?; - // Try to update the lock-file with the new channels - let lock_file = update_lock_file_conda(&project, lock_file, None).await?; + get_up_to_date_prefix(&project, LockFileUsage::Update, args.no_install).await?; project.save()?; - - // Update the installation if needed - if !args.no_install { - // Get the currently installed packages - let prefix = Prefix::new(project.environment_dir())?; - let installed_packages = prefix.find_installed_packages(None).await?; - - // Update the prefix - update_prefix( - project.pypi_package_db()?, - &prefix, - installed_packages, - &lock_file, - Platform::current(), - ) - .await?; - } - // Report back to the user for (name, channel) in missing_channels { eprintln!( diff --git a/src/cli/project/channel/remove.rs b/src/cli/project/channel/remove.rs index 90b9c60e6..205e3d8f4 100644 --- a/src/cli/project/channel/remove.rs +++ b/src/cli/project/channel/remove.rs @@ -1,11 +1,10 @@ -use crate::environment::update_prefix; -use crate::lock_file::{load_lock_file, update_lock_file_conda}; -use crate::prefix::Prefix; +use crate::environment::{get_up_to_date_prefix, LockFileUsage}; + use crate::Project; use clap::Parser; use itertools::Itertools; use miette::IntoDiagnostic; -use rattler_conda_types::{Channel, ChannelConfig, Platform}; +use rattler_conda_types::{Channel, ChannelConfig}; #[derive(Parser, Debug, Default)] pub struct Args { @@ -43,35 +42,15 @@ pub async fn execute(mut project: Project, args: Args) -> miette::Result<()> { return Ok(()); } - // Load the existing lock-file - let lock_file = load_lock_file(&project).await?; - // Remove the channels from the manifest project .manifest .remove_channels(channels_to_remove.iter().map(|(name, _channel)| name))?; // Try to update the lock-file without the removed channels - let lock_file = update_lock_file_conda(&project, lock_file, None).await?; + get_up_to_date_prefix(&project, LockFileUsage::Update, args.no_install).await?; project.save()?; - // Update the installation if needed - if !args.no_install { - // Get the currently installed packages - let prefix = Prefix::new(project.environment_dir())?; - let installed_packages = prefix.find_installed_packages(None).await?; - - // Update the prefix - update_prefix( - project.pypi_package_db()?, - &prefix, - installed_packages, - &lock_file, - Platform::current(), - ) - .await?; - } - // Report back to the user for (name, channel) in channels_to_remove { eprintln!( diff --git a/src/cli/project/platform/add.rs b/src/cli/project/platform/add.rs index db0df24fc..5c9eba918 100644 --- a/src/cli/project/platform/add.rs +++ b/src/cli/project/platform/add.rs @@ -1,8 +1,6 @@ use std::str::FromStr; -use crate::environment::update_prefix; -use crate::lock_file::{load_lock_file, update_lock_file_conda}; -use crate::prefix::Prefix; +use crate::environment::{get_up_to_date_prefix, LockFileUsage}; use crate::Project; use clap::Parser; use itertools::Itertools; @@ -42,33 +40,13 @@ pub async fn execute(mut project: Project, args: Args) -> miette::Result<()> { return Ok(()); } - // Load the existing lock-file - let lock_file = load_lock_file(&project).await?; - // Add the platforms to the lock-file project.manifest.add_platforms(missing_platforms.iter())?; // Try to update the lock-file with the new channels - let lock_file = update_lock_file_conda(&project, lock_file, None).await?; + get_up_to_date_prefix(&project, LockFileUsage::Update, args.no_install).await?; project.save()?; - // Update the installation if needed - if !args.no_install { - // Get the currently installed packages - let prefix = Prefix::new(project.environment_dir())?; - let installed_packages = prefix.find_installed_packages(None).await?; - - // Update the prefix - update_prefix( - project.pypi_package_db()?, - &prefix, - installed_packages, - &lock_file, - Platform::current(), - ) - .await?; - } - // Report back to the user for platform in missing_platforms { eprintln!( diff --git a/src/cli/project/platform/remove.rs b/src/cli/project/platform/remove.rs index 5fb7de9af..9db3726b9 100644 --- a/src/cli/project/platform/remove.rs +++ b/src/cli/project/platform/remove.rs @@ -1,6 +1,5 @@ -use crate::environment::update_prefix; -use crate::lock_file::{load_lock_file, update_lock_file_conda}; -use crate::prefix::Prefix; +use crate::environment::{get_up_to_date_prefix, LockFileUsage}; + use crate::Project; use clap::Parser; use itertools::Itertools; @@ -41,35 +40,14 @@ pub async fn execute(mut project: Project, args: Args) -> miette::Result<()> { return Ok(()); } - // Load the existing lock-file - let lock_file = load_lock_file(&project).await?; - // Remove the platform(s) from the manifest project .manifest .remove_platforms(platforms_to_remove.iter().map(|p| p.to_string()))?; - // Try to update the lock-file without the removed platform(s) - let lock_file = update_lock_file_conda(&project, lock_file, None).await?; + get_up_to_date_prefix(&project, LockFileUsage::Update, args.no_install).await?; project.save()?; - // Update the installation if needed - if !args.no_install { - // Get the currently installed packages - let prefix = Prefix::new(project.environment_dir())?; - let installed_packages = prefix.find_installed_packages(None).await?; - - // Update the prefix - update_prefix( - project.pypi_package_db()?, - &prefix, - installed_packages, - &lock_file, - Platform::current(), - ) - .await?; - } - // Report back to the user for platform in platforms_to_remove { eprintln!( diff --git a/src/cli/remove.rs b/src/cli/remove.rs index 56d0d209b..e4ea1e910 100644 --- a/src/cli/remove.rs +++ b/src/cli/remove.rs @@ -57,7 +57,7 @@ pub async fn execute(args: Args) -> miette::Result<()> { project.save()?; // updating prefix after removing from toml - let _ = get_up_to_date_prefix(&project, LockFileUsage::Update).await?; + let _ = get_up_to_date_prefix(&project, LockFileUsage::Update, false).await?; for (removed, spec) in results.iter().flatten() { let table_name = if let Some(p) = &args.platform { diff --git a/src/cli/run.rs b/src/cli/run.rs index 5b4fcee3c..6c9691419 100644 --- a/src/cli/run.rs +++ b/src/cli/run.rs @@ -163,7 +163,7 @@ pub async fn get_task_env( lock_file_usage: LockFileUsage, ) -> miette::Result> { // Get the prefix which we can then activate. - let prefix = get_up_to_date_prefix(project, lock_file_usage).await?; + let prefix = get_up_to_date_prefix(project, lock_file_usage, false).await?; // Get environment variables from the activation let activation_env = run_activation_async(project, prefix).await?; diff --git a/src/cli/shell.rs b/src/cli/shell.rs index d0238d3d7..c1b914207 100644 --- a/src/cli/shell.rs +++ b/src/cli/shell.rs @@ -203,7 +203,7 @@ pub async fn get_shell_env( lock_file_usage: LockFileUsage, ) -> miette::Result> { // Get the prefix which we can then activate. - let prefix = get_up_to_date_prefix(project, lock_file_usage).await?; + let prefix = get_up_to_date_prefix(project, lock_file_usage, false).await?; // Get environment variables from the activation let activation_env = run_activation_async(project, prefix).await?; diff --git a/src/environment.rs b/src/environment.rs index 1bc2832db..e8f4441b3 100644 --- a/src/environment.rs +++ b/src/environment.rs @@ -9,6 +9,8 @@ use rattler::install::Transaction; use rattler_conda_types::{Platform, PrefixRecord}; use rattler_lock::CondaLock; use rip::index::PackageDb; + +use std::path::PathBuf; use std::{io::ErrorKind, path::Path}; /// Verify the location of the prefix folder is not changed so the applied prefix path is still valid. @@ -99,6 +101,7 @@ pub enum LockFileUsage { pub async fn get_up_to_date_prefix( project: &Project, usage: LockFileUsage, + no_install: bool, ) -> miette::Result { // Make sure the project is in a sane state sanity_check_project(project)?; @@ -117,46 +120,83 @@ pub async fn get_up_to_date_prefix( let mut lock_file = lock_file::load_lock_file(project).await?; let up_to_date = lock_file_satisfies_project(project, &lock_file)?; + if matches!(usage, LockFileUsage::Locked) && !up_to_date { + miette::bail!("Lockfile not up-to-date with the project"); + } + let update_lock = matches!(usage, LockFileUsage::Update) && !up_to_date; - match usage { - LockFileUsage::Update => { - if !up_to_date { - lock_file = lock_file::update_lock_file_conda(project, lock_file, None).await?; - if project.has_pypi_dependencies() { - lock_file = lock_file::update_lock_file_for_pypi(project, lock_file).await?; - } - } + // First lock and install the conda environment + // After which we should have a usable prefix to use for pypi resolution. + if update_lock { + lock_file = lock_file::update_lock_file_conda(project, lock_file, None).await?; + } + + let python_status = if !no_install { + update_prefix_conda( + &prefix, + installed_packages_future.await.into_diagnostic()??, + &lock_file, + Platform::current(), + ) + .await? + } else { + // We don't know and it won't matter because we won't install pypi either + PythonStatus::DoesNotExist + }; + + if project.has_pypi_dependencies() { + if update_lock { + lock_file = lock_file::update_lock_file_for_pypi(project, lock_file).await?; } - LockFileUsage::Locked => { - if !up_to_date { - miette::bail!("Lockfile not up-to-date with the project"); - } + + if !no_install { + // Then update the pypi packages. + update_prefix_pypi( + &prefix, + Platform::current(), + project.pypi_package_db()?, + &lock_file, + &python_status, + ) + .await?; } - // Dont update the lock-file, dont check it - LockFileUsage::Frozen => {} } - // Update the environment - update_prefix( - project.pypi_package_db()?, - &prefix, - installed_packages_future.await.into_diagnostic()??, - &lock_file, - Platform::current(), + Ok(prefix) +} + +pub async fn update_prefix_pypi( + prefix: &Prefix, + platform: Platform, + package_db: &PackageDb, + lock_file: &CondaLock, + status: &PythonStatus, +) -> miette::Result<()> { + // Remove python packages from a previous python distribution if the python version changed. + install_pypi::remove_old_python_distributions(prefix, platform, status)?; + + // Install and/or remove python packages + progress::await_in_progress( + "updating python packages", + install_pypi::update_python_distributions(package_db, prefix, lock_file, platform, status), ) - .await?; + .await +} - Ok(prefix) +#[derive(Clone)] +pub enum PythonStatus { + Changed((u32, u32, u32), PathBuf), + Unchanged((u32, u32, u32), PathBuf), + DoesNotExist, } /// Updates the environment to contain the packages from the specified lock-file -pub async fn update_prefix( - package_db: &PackageDb, +pub async fn update_prefix_conda( prefix: &Prefix, installed_packages: Vec, lock_file: &CondaLock, platform: Platform, -) -> miette::Result<()> { +) -> miette::Result { // Construct a transaction to bring the environment up to date with the lock-file content let desired_conda_packages = lock_file .get_conda_packages_by_platform(platform) @@ -181,22 +221,6 @@ pub async fn update_prefix( .await?; } - // Remove python packages from a previous python distribution if the python version changed. - install_pypi::remove_old_python_distributions(prefix, platform, &transaction)?; - - // Install and/or remove python packages - progress::await_in_progress( - "updating python packages", - install_pypi::update_python_distributions( - package_db, - prefix, - lock_file, - platform, - &transaction, - ), - ) - .await?; - // Mark the location of the prefix create_prefix_location_file( &prefix @@ -207,5 +231,26 @@ pub async fn update_prefix( ) .with_context(|| "failed to create prefix location file.".to_string())?; - Ok(()) + // Determine if the python version changed. + let python_changed = { + if let Some(python_info) = transaction.current_python_info.as_ref() { + let python_version = ( + python_info.short_version.0 as u32, + python_info.short_version.1 as u32, + 0, + ); + + if Some(python_info.short_version) + == transaction.python_info.as_ref().map(|p| p.short_version) + { + PythonStatus::Unchanged(python_version, python_info.path.clone()) + } else { + // Determine the current python distributions in its install locations + PythonStatus::Changed(python_version, python_info.path.clone()) + } + } else { + PythonStatus::DoesNotExist + } + }; + Ok(python_changed) } diff --git a/src/install_pypi.rs b/src/install_pypi.rs index 73d3a27c9..7d6809902 100644 --- a/src/install_pypi.rs +++ b/src/install_pypi.rs @@ -1,3 +1,4 @@ +use crate::environment::PythonStatus; use crate::prefix::Prefix; use crate::progress; use crate::progress::ProgressBarMessageFormatter; @@ -6,8 +7,8 @@ use indexmap::IndexSet; use indicatif::ProgressBar; use itertools::Itertools; use miette::{IntoDiagnostic, WrapErr}; -use rattler::install::Transaction; -use rattler_conda_types::{Platform, PrefixRecord, RepoDataRecord}; + +use rattler_conda_types::Platform; use rattler_lock::{CondaLock, LockedDependency}; use rip::artifacts::wheel::{InstallPaths, UnpackWheelOptions}; use rip::artifacts::Wheel; @@ -32,22 +33,22 @@ pub async fn update_python_distributions( prefix: &Prefix, lock_file: &CondaLock, platform: Platform, - transaction: &Transaction, + status: &PythonStatus, ) -> miette::Result<()> { // Get the python info from the transaction - let Some(python_info) = transaction.python_info.as_ref() else { + let PythonStatus::DoesNotExist = status else { return Ok(()); }; + let (version, path) = match status { + PythonStatus::Changed(version, path) | PythonStatus::Unchanged(version, path) => { + (*version, path) + } + PythonStatus::DoesNotExist => return Ok(()), + }; + // Determine where packages would have been installed - let install_paths = InstallPaths::for_venv( - ( - python_info.short_version.0 as u32, - python_info.short_version.1 as u32, - 0, - ), - platform.is_windows(), - ); + let install_paths = InstallPaths::for_venv(version, platform.is_windows()); // Determine the current python distributions in those locations let current_python_packages = find_distributions_in_venv(prefix.root(), &install_paths) @@ -89,7 +90,7 @@ pub async fn update_python_distributions( let package_install_pb = install_python_distributions( prefix, install_paths, - &prefix.root().join(&python_info.path), + &prefix.root().join(path), package_stream, ) .await?; @@ -297,24 +298,16 @@ fn stream_python_artifacts<'a>( pub fn remove_old_python_distributions( prefix: &Prefix, platform: Platform, - transaction: &Transaction, + python_changed: &PythonStatus, ) -> miette::Result<()> { - // Determine if the current distribution is the same as the desired distribution. - let Some(previous_python_installation) = transaction.current_python_info.as_ref() else { - return Ok(()); + // If the python version didn't change, there is nothing to do here. + let python_version = match python_changed { + PythonStatus::Changed(python_version, _) => *python_version, + PythonStatus::Unchanged(_, _) | PythonStatus::DoesNotExist => { + return Ok(()); + } }; - if Some(previous_python_installation.short_version) - == transaction.python_info.as_ref().map(|p| p.short_version) - { - return Ok(()); - } - // Determine the current python distributions in its install locations - let python_version = ( - previous_python_installation.short_version.0 as u32, - previous_python_installation.short_version.1 as u32, - 0, - ); let install_paths = InstallPaths::for_venv(python_version, platform.is_windows()); // Locate the packages that are installed in the previous environment diff --git a/src/lock_file/mod.rs b/src/lock_file/mod.rs index fc983fd02..f2aae6683 100644 --- a/src/lock_file/mod.rs +++ b/src/lock_file/mod.rs @@ -157,7 +157,7 @@ pub async fn update_lock_file_conda( pub async fn update_lock_file_for_pypi( project: &Project, - new_lock_file: CondaLock, + lock_for_conda: CondaLock, ) -> miette::Result { let platforms = project.platforms(); let _top_level_progress = @@ -166,7 +166,7 @@ pub async fn update_lock_file_for_pypi( let records = platforms .iter() - .map(|plat| new_lock_file.get_conda_packages_by_platform(*plat)); + .map(|plat| lock_for_conda.get_conda_packages_by_platform(*plat)); let result: miette::Result> = stream::iter(izip!(platforms.iter(), solve_bars.iter().cloned(), records)) @@ -225,16 +225,23 @@ pub async fn update_lock_file_for_pypi( let conda_lock = builder.build().into_diagnostic()?; // TODO: think of a better way to do this - Ok(CondaLock { - metadata: new_lock_file.metadata, + // Seeing as we are not using the content-hash anyways this seems to be fine + let latest_lock = CondaLock { + metadata: lock_for_conda.metadata, package: Vec::from_iter( conda_lock .package .iter() - .chain(new_lock_file.package.iter()) + .chain(lock_for_conda.package.iter()) .cloned(), ), - }) + }; + + // Write the conda lock to disk + conda_lock + .to_path(&project.lock_file_path()) + .into_diagnostic()?; + Ok(latest_lock) } async fn resolve_pypi( From 018bb34895c3f24ff10224ff9f03e911d8d49bfd Mon Sep 17 00:00:00 2001 From: Tim de Jager Date: Mon, 1 Jan 2024 21:19:09 +0100 Subject: [PATCH 05/10] refactor: everything uses `get_up_to_date_prefix` --- src/cli/add.rs | 75 +++++++++++--------------------------------------- 1 file changed, 16 insertions(+), 59 deletions(-) diff --git a/src/cli/add.rs b/src/cli/add.rs index 21939e9f4..96e07e794 100644 --- a/src/cli/add.rs +++ b/src/cli/add.rs @@ -1,17 +1,14 @@ -use crate::environment::{update_prefix_conda, verify_prefix_location_unchanged}; -use crate::lock_file::update_lock_file_for_pypi; -use crate::prefix::Prefix; +use crate::environment::{get_up_to_date_prefix, verify_prefix_location_unchanged, LockFileUsage}; + use crate::project::{DependencyType, SpecType}; use crate::{ - consts, - lock_file::{load_lock_file, update_lock_file_conda}, - project::python::PyPiRequirement, - project::Project, + consts, project::python::PyPiRequirement, project::Project, virtual_packages::get_minimal_virtual_packages, }; use clap::Parser; use indexmap::IndexMap; use itertools::Itertools; + use miette::{IntoDiagnostic, WrapErr}; use rattler_conda_types::version_spec::{LogicalOperator, RangeOperator}; use rattler_conda_types::{ @@ -219,8 +216,13 @@ pub async fn add_pypi_specs_to_project( } } } + let lock_file_usage = if no_update_lockfile { + LockFileUsage::Frozen + } else { + LockFileUsage::Update + }; - update_environment(project, None, no_install, no_update_lockfile).await?; + get_up_to_date_prefix(project, lock_file_usage, no_install).await?; project.save()?; @@ -318,62 +320,17 @@ pub async fn add_conda_specs_to_project( } } } - project.save()?; - - update_environment( - project, - Some(sparse_repo_data), - no_install, - no_update_lockfile, - ) - .await?; - - Ok(()) -} - -/// Updates the lock file and potentially the prefix to get an up-to-date environment. -/// -/// We are using this function instead of [`crate::environment::get_up_to_date_prefix`] because we want to be able to -/// specify if we do not want to update the prefix. Also we know the lock file needs to be updated so `--frozen` and `--locked` -/// make no sense in this scenario. -/// -/// Essentially, other than that it does almost the same thing -async fn update_environment( - project: &Project, - sparse_repo_data: Option>, - no_install: bool, - no_update_lockfile: bool, -) -> miette::Result<()> { - // Update the lock file - let lock_file = if !no_update_lockfile { - let lock = - update_lock_file_conda(project, load_lock_file(project).await?, sparse_repo_data) - .await?; - - if project.has_pypi_dependencies() { - update_lock_file_for_pypi(project, lock).await?.into() - } else { - lock.into() - } + let lock_file_usage = if no_update_lockfile { + LockFileUsage::Frozen } else { - None + LockFileUsage::Update }; + get_up_to_date_prefix(project, lock_file_usage, no_install).await?; + project.save()?; - if let Some(lock_file) = lock_file { - if !no_install { - crate::environment::sanity_check_project(project)?; - - // Get the currently installed packages - let prefix = Prefix::new(project.environment_dir())?; - let installed_packages = prefix.find_installed_packages(None).await?; - - // Update the prefix - update_prefix_conda(&prefix, installed_packages, &lock_file, Platform::current()) - .await?; - } - } Ok(()) } + /// Given several specs determines the highest installable version for them. pub fn determine_best_version( new_specs: &HashMap, From 77e47722b0f76dce8ebd41ba1a04089b9ed449a2 Mon Sep 17 00:00:00 2001 From: Bas Zalmstra Date: Tue, 2 Jan 2024 16:55:55 +0100 Subject: [PATCH 06/10] split usage a little --- src/environment.rs | 111 +++++++++++++++++++++++++++++--------------- src/install_pypi.rs | 37 +++++++++------ 2 files changed, 97 insertions(+), 51 deletions(-) diff --git a/src/environment.rs b/src/environment.rs index 9abb91d54..43e9bad49 100644 --- a/src/environment.rs +++ b/src/environment.rs @@ -5,11 +5,10 @@ use crate::{ use miette::{Context, IntoDiagnostic, LabeledSpan}; use crate::lock_file::lock_file_satisfies_project; -use rattler::install::Transaction; -use rattler_conda_types::{Platform, PrefixRecord}; +use rattler::install::{PythonInfo, Transaction}; +use rattler_conda_types::{Platform, PrefixRecord, RepoDataRecord}; use rattler_lock::CondaLock; use rip::index::PackageDb; -use std::path::PathBuf; use std::{io::ErrorKind, path::Path}; /// Verify the location of the prefix folder is not changed so the applied prefix path is still valid. @@ -84,7 +83,7 @@ pub fn sanity_check_project(project: &Project) -> miette::Result<()> { } /// Specifies how the lock-file should be updated. -#[derive(Debug, Default)] +#[derive(Debug, Default, PartialEq, Eq, Copy, Clone)] pub enum LockFileUsage { /// Update the lock-file if it is out of date. #[default] @@ -95,6 +94,24 @@ pub enum LockFileUsage { Frozen, } +impl LockFileUsage { + /// Returns true if the lock-file should be updated if it is out of date. + pub fn allows_lock_file_updates(self) -> bool { + match self { + LockFileUsage::Update => true, + LockFileUsage::Locked | LockFileUsage::Frozen => false, + } + } + + /// Returns true if the lock-file should be checked if it is out of date. + pub fn should_check_if_out_of_date(self) -> bool { + match self { + LockFileUsage::Update | LockFileUsage::Locked => true, + LockFileUsage::Frozen => false, + } + } +} + /// Returns the prefix associated with the given environment. If the prefix doesn't exist or is not /// up to date it is updated. pub async fn get_up_to_date_prefix( @@ -112,21 +129,29 @@ pub async fn get_up_to_date_prefix( tokio::spawn(async move { prefix.find_installed_packages(None).await }) }; - // Update the lock-file if it is out of date. - if matches!(usage, LockFileUsage::Frozen) && !project.lock_file_path().is_file() { - miette::bail!("No lockfile available, can't do a frozen installation."); + // If there is no lock-file and we are also not allowed to update it, we can bail immediately. + if !project.lock_file_path().is_file() && !usage.allows_lock_file_updates() { + miette::bail!("no lockfile available, can't do a frozen installation."); } + // Load the lock-file into memory. let mut lock_file = lock_file::load_lock_file(project).await?; - let up_to_date = lock_file_satisfies_project(project, &lock_file)?; - if matches!(usage, LockFileUsage::Locked) && !up_to_date { - miette::bail!("Lockfile not up-to-date with the project"); - } - let update_lock = matches!(usage, LockFileUsage::Update) && !up_to_date; + + // Check if the lock-file is up to date, but only if the current usage allows it. + let update_lock_file = if usage.should_check_if_out_of_date() + && !lock_file_satisfies_project(project, &lock_file)? + { + if !usage.allows_lock_file_updates() { + miette::bail!("lockfile not up-to-date with the project"); + } + true + } else { + false + }; // First lock and install the conda environment // After which we should have a usable prefix to use for pypi resolution. - if update_lock { + if update_lock_file { lock_file = lock_file::update_lock_file_conda(project, lock_file, None).await?; } @@ -144,7 +169,7 @@ pub async fn get_up_to_date_prefix( }; if project.has_pypi_dependencies() { - if update_lock { + if update_lock_file { lock_file = lock_file::update_lock_file_for_pypi(project, lock_file).await?; } @@ -184,11 +209,43 @@ pub async fn update_prefix_pypi( #[derive(Clone)] pub enum PythonStatus { - Changed((u32, u32, u32), PathBuf), - Unchanged((u32, u32, u32), PathBuf), + /// The python interpreter changed from `old` to `new`. + Changed { old: PythonInfo, new: PythonInfo }, + + /// The python interpreter remained the same. + Unchanged(PythonInfo), + + /// The python interpreter was removed from the environment + Removed { old: PythonInfo }, + + /// The python interpreter was added to the environment + Added { new: PythonInfo }, + + /// There is no python interpreter in the environment. DoesNotExist, } +impl PythonStatus { + /// Determine the [`PythonStatus`] from a [`Transaction`]. + pub fn from_transaction(transaction: &Transaction) -> Self { + match ( + transaction.current_python_info.as_ref(), + transaction.python_info.as_ref(), + ) { + (Some(old), Some(new)) if old.short_version != new.short_version => { + PythonStatus::Changed { + old: old.clone(), + new: new.clone(), + } + } + (Some(_), Some(new)) => PythonStatus::Unchanged(new.clone()), + (None, Some(new)) => PythonStatus::Added { new: new.clone() }, + (Some(old), None) => PythonStatus::Removed { old: old.clone() }, + (None, None) => PythonStatus::DoesNotExist, + } + } +} + /// Updates the environment to contain the packages from the specified lock-file pub async fn update_prefix_conda( prefix: &Prefix, @@ -231,25 +288,5 @@ pub async fn update_prefix_conda( .with_context(|| "failed to create prefix location file.".to_string())?; // Determine if the python version changed. - let python_changed = { - if let Some(python_info) = transaction.current_python_info.as_ref() { - let python_version = ( - python_info.short_version.0 as u32, - python_info.short_version.1 as u32, - 0, - ); - - if Some(python_info.short_version) - == transaction.python_info.as_ref().map(|p| p.short_version) - { - PythonStatus::Unchanged(python_version, python_info.path.clone()) - } else { - // Determine the current python distributions in its install locations - PythonStatus::Changed(python_version, python_info.path.clone()) - } - } else { - PythonStatus::DoesNotExist - } - }; - Ok(python_changed) + Ok(PythonStatus::from_transaction(&transaction)) } diff --git a/src/install_pypi.rs b/src/install_pypi.rs index 7d6809902..6d3c6beb3 100644 --- a/src/install_pypi.rs +++ b/src/install_pypi.rs @@ -35,20 +35,23 @@ pub async fn update_python_distributions( platform: Platform, status: &PythonStatus, ) -> miette::Result<()> { - // Get the python info from the transaction - let PythonStatus::DoesNotExist = status else { - return Ok(()); - }; - - let (version, path) = match status { - PythonStatus::Changed(version, path) | PythonStatus::Unchanged(version, path) => { - (*version, path) + let python_info = match status { + PythonStatus::Changed { new, .. } + | PythonStatus::Unchanged(new) + | PythonStatus::Added { new } => new, + PythonStatus::Removed { .. } | PythonStatus::DoesNotExist => { + // No python interpreter in the environment, so there is nothing to do here. + return Ok(()); } - PythonStatus::DoesNotExist => return Ok(()), }; // Determine where packages would have been installed - let install_paths = InstallPaths::for_venv(version, platform.is_windows()); + let python_version = ( + python_info.short_version.0 as u32, + python_info.short_version.1 as u32, + 0, + ); + let install_paths = InstallPaths::for_venv(python_version, platform.is_windows()); // Determine the current python distributions in those locations let current_python_packages = find_distributions_in_venv(prefix.root(), &install_paths) @@ -90,7 +93,7 @@ pub async fn update_python_distributions( let package_install_pb = install_python_distributions( prefix, install_paths, - &prefix.root().join(path), + &prefix.root().join(python_info.path()), package_stream, ) .await?; @@ -302,12 +305,18 @@ pub fn remove_old_python_distributions( ) -> miette::Result<()> { // If the python version didn't change, there is nothing to do here. let python_version = match python_changed { - PythonStatus::Changed(python_version, _) => *python_version, - PythonStatus::Unchanged(_, _) | PythonStatus::DoesNotExist => { - return Ok(()); + PythonStatus::Removed { old } | PythonStatus::Changed { old, .. } => old, + PythonStatus::Added { .. } | PythonStatus::DoesNotExist | PythonStatus::Unchanged(_) => { + return Ok(()) } }; + // Get the interpreter version from the info + let python_version = ( + python_version.short_version.0 as u32, + python_version.short_version.1 as u32, + 0, + ); let install_paths = InstallPaths::for_venv(python_version, platform.is_windows()); // Locate the packages that are installed in the previous environment From bbae3c22e62d73bb059c013c661f57d2a1c39877 Mon Sep 17 00:00:00 2001 From: Bas Zalmstra Date: Thu, 4 Jan 2024 09:45:06 +0100 Subject: [PATCH 07/10] fix: pass repodata down to lock-file updates --- src/cli/add.rs | 4 ++-- src/cli/install.rs | 2 +- src/cli/project/channel/add.rs | 2 +- src/cli/project/channel/remove.rs | 2 +- src/cli/project/platform/add.rs | 2 +- src/cli/project/platform/remove.rs | 2 +- src/cli/remove.rs | 2 +- src/cli/run.rs | 2 +- src/cli/shell.rs | 2 +- src/environment.rs | 9 ++++++++- 10 files changed, 18 insertions(+), 11 deletions(-) diff --git a/src/cli/add.rs b/src/cli/add.rs index 96e07e794..2f5a1a671 100644 --- a/src/cli/add.rs +++ b/src/cli/add.rs @@ -222,7 +222,7 @@ pub async fn add_pypi_specs_to_project( LockFileUsage::Update }; - get_up_to_date_prefix(project, lock_file_usage, no_install).await?; + get_up_to_date_prefix(project, lock_file_usage, no_install, None).await?; project.save()?; @@ -325,7 +325,7 @@ pub async fn add_conda_specs_to_project( } else { LockFileUsage::Update }; - get_up_to_date_prefix(project, lock_file_usage, no_install).await?; + get_up_to_date_prefix(project, lock_file_usage, no_install, Some(sparse_repo_data)).await?; project.save()?; Ok(()) diff --git a/src/cli/install.rs b/src/cli/install.rs index ac475b911..f72c4f3ab 100644 --- a/src/cli/install.rs +++ b/src/cli/install.rs @@ -17,7 +17,7 @@ pub struct Args { pub async fn execute(args: Args) -> miette::Result<()> { let project = Project::load_or_else_discover(args.manifest_path.as_deref())?; - get_up_to_date_prefix(&project, args.lock_file_usage.into(), false).await?; + get_up_to_date_prefix(&project, args.lock_file_usage.into(), false, None).await?; // Emit success eprintln!( diff --git a/src/cli/project/channel/add.rs b/src/cli/project/channel/add.rs index 40aed24f6..7dd22f9f0 100644 --- a/src/cli/project/channel/add.rs +++ b/src/cli/project/channel/add.rs @@ -51,7 +51,7 @@ pub async fn execute(mut project: Project, args: Args) -> miette::Result<()> { .manifest .add_channels(missing_channels.iter().map(|(name, _channel)| name))?; - get_up_to_date_prefix(&project, LockFileUsage::Update, args.no_install).await?; + get_up_to_date_prefix(&project, LockFileUsage::Update, args.no_install, None).await?; project.save()?; // Report back to the user for (name, channel) in missing_channels { diff --git a/src/cli/project/channel/remove.rs b/src/cli/project/channel/remove.rs index 205e3d8f4..6fd5fa06d 100644 --- a/src/cli/project/channel/remove.rs +++ b/src/cli/project/channel/remove.rs @@ -48,7 +48,7 @@ pub async fn execute(mut project: Project, args: Args) -> miette::Result<()> { .remove_channels(channels_to_remove.iter().map(|(name, _channel)| name))?; // Try to update the lock-file without the removed channels - get_up_to_date_prefix(&project, LockFileUsage::Update, args.no_install).await?; + get_up_to_date_prefix(&project, LockFileUsage::Update, args.no_install, None).await?; project.save()?; // Report back to the user diff --git a/src/cli/project/platform/add.rs b/src/cli/project/platform/add.rs index 5c9eba918..dca4f2ab3 100644 --- a/src/cli/project/platform/add.rs +++ b/src/cli/project/platform/add.rs @@ -44,7 +44,7 @@ pub async fn execute(mut project: Project, args: Args) -> miette::Result<()> { project.manifest.add_platforms(missing_platforms.iter())?; // Try to update the lock-file with the new channels - get_up_to_date_prefix(&project, LockFileUsage::Update, args.no_install).await?; + get_up_to_date_prefix(&project, LockFileUsage::Update, args.no_install, None).await?; project.save()?; // Report back to the user diff --git a/src/cli/project/platform/remove.rs b/src/cli/project/platform/remove.rs index 9db3726b9..97c4fc37a 100644 --- a/src/cli/project/platform/remove.rs +++ b/src/cli/project/platform/remove.rs @@ -45,7 +45,7 @@ pub async fn execute(mut project: Project, args: Args) -> miette::Result<()> { .manifest .remove_platforms(platforms_to_remove.iter().map(|p| p.to_string()))?; - get_up_to_date_prefix(&project, LockFileUsage::Update, args.no_install).await?; + get_up_to_date_prefix(&project, LockFileUsage::Update, args.no_install, None).await?; project.save()?; // Report back to the user diff --git a/src/cli/remove.rs b/src/cli/remove.rs index e4ea1e910..6da6eadc9 100644 --- a/src/cli/remove.rs +++ b/src/cli/remove.rs @@ -57,7 +57,7 @@ pub async fn execute(args: Args) -> miette::Result<()> { project.save()?; // updating prefix after removing from toml - let _ = get_up_to_date_prefix(&project, LockFileUsage::Update, false).await?; + let _ = get_up_to_date_prefix(&project, LockFileUsage::Update, false, None).await?; for (removed, spec) in results.iter().flatten() { let table_name = if let Some(p) = &args.platform { diff --git a/src/cli/run.rs b/src/cli/run.rs index 6c9691419..7405276d5 100644 --- a/src/cli/run.rs +++ b/src/cli/run.rs @@ -163,7 +163,7 @@ pub async fn get_task_env( lock_file_usage: LockFileUsage, ) -> miette::Result> { // Get the prefix which we can then activate. - let prefix = get_up_to_date_prefix(project, lock_file_usage, false).await?; + let prefix = get_up_to_date_prefix(project, lock_file_usage, false, None).await?; // Get environment variables from the activation let activation_env = run_activation_async(project, prefix).await?; diff --git a/src/cli/shell.rs b/src/cli/shell.rs index c1b914207..2731b5db6 100644 --- a/src/cli/shell.rs +++ b/src/cli/shell.rs @@ -203,7 +203,7 @@ pub async fn get_shell_env( lock_file_usage: LockFileUsage, ) -> miette::Result> { // Get the prefix which we can then activate. - let prefix = get_up_to_date_prefix(project, lock_file_usage, false).await?; + let prefix = get_up_to_date_prefix(project, lock_file_usage, false, None).await?; // Get environment variables from the activation let activation_env = run_activation_async(project, prefix).await?; diff --git a/src/environment.rs b/src/environment.rs index 43e9bad49..ad76e8b1f 100644 --- a/src/environment.rs +++ b/src/environment.rs @@ -8,6 +8,7 @@ use crate::lock_file::lock_file_satisfies_project; use rattler::install::{PythonInfo, Transaction}; use rattler_conda_types::{Platform, PrefixRecord, RepoDataRecord}; use rattler_lock::CondaLock; +use rattler_repodata_gateway::sparse::SparseRepoData; use rip::index::PackageDb; use std::{io::ErrorKind, path::Path}; @@ -114,10 +115,16 @@ impl LockFileUsage { /// Returns the prefix associated with the given environment. If the prefix doesn't exist or is not /// up to date it is updated. +/// +/// The `sparse_repo_data` is used when the lock-file is update. We pass it into this function to +/// make sure the data is not loaded twice since the repodata takes up a lot of memory and takes a +/// while to load. If `sparse_repo_data` is `None` it will be downloaded. If the lock-file is not +/// updated, the `sparse_repo_data` is ignored. pub async fn get_up_to_date_prefix( project: &Project, usage: LockFileUsage, no_install: bool, + sparse_repo_data: Option>, ) -> miette::Result { // Make sure the project is in a sane state sanity_check_project(project)?; @@ -152,7 +159,7 @@ pub async fn get_up_to_date_prefix( // First lock and install the conda environment // After which we should have a usable prefix to use for pypi resolution. if update_lock_file { - lock_file = lock_file::update_lock_file_conda(project, lock_file, None).await?; + lock_file = lock_file::update_lock_file_conda(project, lock_file, sparse_repo_data).await?; } let python_status = if !no_install { From 6997c543aa694231dd956bc0c2dee42c3c164432 Mon Sep 17 00:00:00 2001 From: Bas Zalmstra Date: Thu, 4 Jan 2024 11:15:01 +0100 Subject: [PATCH 08/10] fix: write out conda packages --- src/lock_file/mod.rs | 24 ++++++++++++++---------- 1 file changed, 14 insertions(+), 10 deletions(-) diff --git a/src/lock_file/mod.rs b/src/lock_file/mod.rs index f2aae6683..092447ef4 100644 --- a/src/lock_file/mod.rs +++ b/src/lock_file/mod.rs @@ -17,7 +17,7 @@ use rattler_lock::{ CondaLockedDependencyBuilder, LockFileBuilder, LockedPackagesBuilder, PypiLockedDependencyBuilder, }, - CondaLock, PackageHashes, + CondaLock, LockedDependencyKind, PackageHashes, }; use rattler_repodata_gateway::sparse::SparseRepoData; use rattler_solve::{resolvo, SolverImpl}; @@ -222,25 +222,29 @@ pub async fn update_lock_file_for_pypi( for locked_packages in result? { builder = builder.add_locked_packages(locked_packages); } - let conda_lock = builder.build().into_diagnostic()?; + let conda_lock_pypi_only = builder.build().into_diagnostic()?; // TODO: think of a better way to do this // Seeing as we are not using the content-hash anyways this seems to be fine let latest_lock = CondaLock { metadata: lock_for_conda.metadata, - package: Vec::from_iter( - conda_lock - .package - .iter() - .chain(lock_for_conda.package.iter()) - .cloned(), - ), + package: conda_lock_pypi_only + .package + .into_iter() + .chain( + lock_for_conda + .package + .into_iter() + .filter(|p| matches!(p.kind, LockedDependencyKind::Conda(_))), + ) + .collect(), }; // Write the conda lock to disk - conda_lock + latest_lock .to_path(&project.lock_file_path()) .into_diagnostic()?; + Ok(latest_lock) } From a2b929e694cd104b480c7c136e500431ed8218f9 Mon Sep 17 00:00:00 2001 From: Bas Zalmstra Date: Thu, 4 Jan 2024 11:20:32 +0100 Subject: [PATCH 09/10] Update src/lock_file/mod.rs Co-authored-by: Ruben Arts --- src/lock_file/mod.rs | 3 --- 1 file changed, 3 deletions(-) diff --git a/src/lock_file/mod.rs b/src/lock_file/mod.rs index 092447ef4..f5cb0bd3d 100644 --- a/src/lock_file/mod.rs +++ b/src/lock_file/mod.rs @@ -352,9 +352,6 @@ async fn resolve_platform( locked_packages.add_locked_package(locked_package); } - // Add pypi packages - // let locked_packages = - // resolve_pypi(project, &mut records, locked_packages, platform, &pb).await?; Ok(locked_packages) } From 8e634b690a9935b7ee1c791fb068f957839890f0 Mon Sep 17 00:00:00 2001 From: Bas Zalmstra Date: Thu, 4 Jan 2024 11:22:24 +0100 Subject: [PATCH 10/10] fix: fmt --- src/lock_file/mod.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/src/lock_file/mod.rs b/src/lock_file/mod.rs index f5cb0bd3d..335b53361 100644 --- a/src/lock_file/mod.rs +++ b/src/lock_file/mod.rs @@ -352,7 +352,6 @@ async fn resolve_platform( locked_packages.add_locked_package(locked_package); } - Ok(locked_packages) }