Split pypi and conda solve steps (#601)

Did a large refactor that does two main things:

* Split the solve, into a conda solve and a pypi solve after. This will
make it easier to use the conda environment and python interpreter when
using sdist resolution.
* I changed the `get_up_to_date_prefix` function so it accepts a
`no_install` flag. This way we can actually use it everywhere instead of
duplicating a lot of code.

Later on we will need to actually use the interpreter in the sdist
resolution step.

PR depends on #589, need to merge that first.

---------

Co-authored-by: Bas Zalmstra <bas@prefix.dev>
Co-authored-by: Bas Zalmstra <zalmstra.bas@gmail.com>
Co-authored-by: Ruben Arts <ruben@prefix.dev>

src/cli/add.rs

@@ -1,25 +1,24 @@
-use crate::environment::{update_prefix, verify_prefix_location_unchanged};
-use crate::prefix::Prefix;
-use crate::project::{DependencyType, SpecType};
 use crate::{
     consts,
-    lock_file::{load_lock_file, update_lock_file},
-    project::python::PyPiRequirement,
-    project::Project,
+    environment::{get_up_to_date_prefix, verify_prefix_location_unchanged, LockFileUsage},
+    project::{python::PyPiRequirement, DependencyType, Project, SpecType},
 };
 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::{
+    version_spec::{LogicalOperator, RangeOperator},
     MatchSpec, NamelessMatchSpec, PackageName, Platform, Version, VersionSpec,
 };
 use rattler_repodata_gateway::sparse::SparseRepoData;
 use rattler_solve::{resolvo, SolverImpl};
-use std::collections::{HashMap, HashSet};
-use std::path::PathBuf;
-use std::str::FromStr;
+use std::{
+    collections::{HashMap, HashSet},
+    path::PathBuf,
+    str::FromStr,
+};
 
 /// Adds a dependency to the project
 #[derive(Parser, Debug, Default)]
@@ -217,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, None).await?;
 
     project.save()?;
 
@@ -317,60 +321,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 {
-        Some(update_lock_file(project, load_lock_file(project).await?, sparse_repo_data).await?)
+    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, Some(sparse_repo_data)).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(
-                project.pypi_package_db()?,
-                &prefix,
-                installed_packages,
-                &lock_file,
-                Platform::current(),
-            )
-            .await?;
-        }
-    }
     Ok(())
 }
+
 /// Given several specs determines the highest installable version for them.
 pub fn determine_best_version(
     project: &Project,

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, None).await?;
 
     // Emit success
     eprintln!(

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};
-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(&project, lock_file, None).await?;
+    get_up_to_date_prefix(&project, LockFileUsage::Update, args.no_install, None).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!(

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};
-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(&project, lock_file, None).await?;
+    get_up_to_date_prefix(&project, LockFileUsage::Update, args.no_install, None).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!(

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};
-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(&project, lock_file, None).await?;
+    get_up_to_date_prefix(&project, LockFileUsage::Update, args.no_install, None).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!(

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};
-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(&project, lock_file, None).await?;
+    get_up_to_date_prefix(&project, LockFileUsage::Update, args.no_install, None).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!(

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, None).await?;
 
     for (removed, spec) in results.iter().flatten() {
         let table_name = if let Some(p) = &args.platform {

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, None).await?;
 
     // Get environment variables from the activation
     let activation_env = run_activation_async(project, prefix).await?;

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, None).await?;
 
     // Get environment variables from the activation
     let activation_env = run_activation_async(project, prefix).await?;