From 4e620f41324f2fb02957fe1aeeb1d58ebae23489 Mon Sep 17 00:00:00 2001
From: Tim de Jager <tim@prefix.dev>
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<PrefixRecord, RepoDataRecord>,
-) -> 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<Item = miette::Result<(Option<String>, HashSet<Extra>, Wheel)>> + Sized,
-) -> miette::Result<Option<ProgressBar>> {
-    // 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<Item = miette::Result<(Option<String>, HashSet<Extra>, Wheel)>> + 'a,
-    Option<ProgressBar>,
-) {
-    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<PrefixRecord, RepoDataRecord>,
-) -> 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<Distribution>,
-    desired_python_packages: Vec<&'p LockedDependency>,
-) -> (Vec<Distribution>, 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<IndexSet<WheelTag>>)> {
-    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<WheelTag>>),
-) -> 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<PrefixRecord, RepoDataRecord>,
+) -> 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<Item = miette::Result<(Option<String>, HashSet<Extra>, Wheel)>> + Sized,
+) -> miette::Result<Option<ProgressBar>> {
+    // 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<Item = miette::Result<(Option<String>, HashSet<Extra>, Wheel)>> + 'a,
+    Option<ProgressBar>,
+) {
+    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<PrefixRecord, RepoDataRecord>,
+) -> 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<Distribution>,
+    desired_python_packages: Vec<&'p LockedDependency>,
+) -> (Vec<Distribution>, 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<IndexSet<WheelTag>>)> {
+    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<WheelTag>>),
+) -> 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 <tim@prefix.dev>
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 <tim@prefix.dev>
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<CondaLock> {
     .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<ProgressBar> {
+    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<Vec<SparseRepoData>>,
@@ -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<Vec<_>> =
         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<CondaLock> {
+    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<Vec<_>> =
+        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<LockedPackagesBuilder> {
+    // 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<Vec<PinnedPackage<'p>>> {
     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 <tim@prefix.dev>
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<HashMap<String, String>> {
     // 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<HashMap<String, String>> {
     // 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<Prefix> {
     // 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<PrefixRecord>,
     lock_file: &CondaLock,
     platform: Platform,
-) -> miette::Result<()> {
+) -> miette::Result<PythonStatus> {
     // 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<PrefixRecord, RepoDataRecord>,
+    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<PrefixRecord, RepoDataRecord>,
+    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<CondaLock> {
     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<Vec<_>> =
         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 <tim@prefix.dev>
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<Vec<SparseRepoData>>,
-    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<PackageName, NamelessMatchSpec>,

From 77e47722b0f76dce8ebd41ba1a04089b9ed449a2 Mon Sep 17 00:00:00 2001
From: Bas Zalmstra <bas@prefix.dev>
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<PrefixRecord, RepoDataRecord>) -> 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 <bas@prefix.dev>
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<HashMap<String, String>> {
     // 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<HashMap<String, String>> {
     // 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<Vec<SparseRepoData>>,
 ) -> miette::Result<Prefix> {
     // 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 <bas@prefix.dev>
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 <zalmstra.bas@gmail.com>
Date: Thu, 4 Jan 2024 11:20:32 +0100
Subject: [PATCH 09/10] Update src/lock_file/mod.rs

Co-authored-by: Ruben Arts <ruben@prefix.dev>
---
 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 <bas@prefix.dev>
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)
 }