refactor: add to use update code (#1508)

Refactors the `pixi add` code to use the same code path as `update` and
`install`. This provides:

* A more correct result
* Speed, because now everything is done in parallel
* Progress Reporting!
* Version caps for pypi deps:
Fixes: #63
Fixes: #1506

---------

Co-authored-by: Tim de Jager <[email protected]>

src/cli/init.rs

@@ -1,6 +1,7 @@
 use crate::config::Config;
 use crate::environment::{get_up_to_date_prefix, LockFileUsage};
 use crate::project::manifest::pyproject::PyProjectToml;
+use crate::project::manifest::DependencyOverwriteBehavior;
 use crate::utils::conda_environment_file::CondaEnvFile;
 use crate::{config::get_default_author, consts};
 use crate::{FeatureName, Project};
@@ -188,6 +189,7 @@ pub async fn execute(args: Args) -> miette::Result<()> {
                 crate::SpecType::Run,
                 &platforms,
                 &FeatureName::default(),
+                DependencyOverwriteBehavior::Overwrite,
             )?;
         }
         for requirement in pypi_deps {
@@ -196,6 +198,7 @@ pub async fn execute(args: Args) -> miette::Result<()> {
                 &platforms,
                 &FeatureName::default(),
                 None,
+                DependencyOverwriteBehavior::Overwrite,
             )?;
         }
         project.save()?;

src/cli/remove.rs

@@ -63,7 +63,7 @@ pub async fn execute(args: Args) -> miette::Result<()> {
     )
     .await?;
 
-    args.display_success("Removed");
+    args.display_success("Removed", Default::default());
 
     Project::warn_on_discovered_from_env(args.manifest_path.as_deref());
     Ok(())

src/cli/update.rs

@@ -12,7 +12,7 @@ use indexmap::IndexMap;
 use itertools::{Either, Itertools};
 use miette::{Context, IntoDiagnostic, MietteDiagnostic};
 use rattler_conda_types::Platform;
-use rattler_lock::{LockFile, LockFileBuilder, Package};
+use rattler_lock::{LockFile, Package};
 use serde::Serialize;
 use serde_json::Value;
 use tabwriter::TabWriter;
@@ -22,8 +22,7 @@ use crate::{
     consts,
     consts::{CondaEmoji, PypiEmoji},
     load_lock_file,
-    lock_file::UpdateContext,
-    project::grouped_environment::GroupedEnvironment,
+    lock_file::{filter_lock_file, UpdateContext},
     EnvironmentName, HasFeatures, Project,
 };
 
@@ -262,35 +261,9 @@ fn check_package_exists(
 
 /// Constructs a new lock-file where some of the constraints have been removed.
 fn unlock_packages(project: &Project, lock_file: &LockFile, specs: &UpdateSpecs) -> LockFile {
-    let mut builder = LockFileBuilder::new();
-
-    for (environment_name, environment) in lock_file.environments() {
-        // Find the environment in the project
-        let Some(project_env) = project.environment(environment_name) else {
-            continue;
-        };
-
-        // Copy the channels
-        builder.set_channels(environment_name, environment.channels().to_vec());
-
-        // Copy the indexes
-        let indexes = environment
-            .pypi_indexes()
-            .cloned()
-            .unwrap_or_else(|| GroupedEnvironment::from(project_env).pypi_options().into());
-        builder.set_pypi_indexes(environment_name, indexes);
-
-        // Copy all packages that don't need to be relaxed
-        for (platform, packages) in environment.packages_by_platform() {
-            for package in packages {
-                if !specs.should_relax(environment_name, platform, &package) {
-                    builder.add_package(environment_name, platform, package);
-                }
-            }
-        }
-    }
-
-    builder.finish()
+    filter_lock_file(project, lock_file, |env, platform, package| {
+        !specs.should_relax(env.name().as_str(), platform, package)
+    })
 }
 
 // Represents the differences between two sets of packages.

src/lock_file/mod.rs

@@ -4,32 +4,35 @@ mod records_by_name;
 mod resolve;
 mod satisfiability;
 mod update;
+mod utils;
 
-use crate::Project;
 use miette::{IntoDiagnostic, WrapErr};
-use rattler_conda_types::RepoDataRecord;
-use rattler_lock::{LockFile, PypiPackageData, PypiPackageEnvironmentData};
-
 pub use outdated::OutdatedEnvironments;
 pub use package_identifier::PypiPackageIdentifier;
+use rattler_conda_types::RepoDataRecord;
+use rattler_lock::{LockFile, PypiPackageData, PypiPackageEnvironmentData};
 pub use records_by_name::{PypiRecordsByName, RepoDataRecordsByName};
 pub use resolve::{
     conda::resolve_conda, pypi::resolve_pypi, uv_resolution_context::UvResolutionContext,
 };
 pub use satisfiability::{verify_environment_satisfiability, verify_platform_satisfiability};
 pub use update::{LockFileDerivedData, UpdateContext, UpdateLockFileOptions};
+pub use utils::filter_lock_file;
+
+use crate::Project;
 
 /// A list of conda packages that are locked for a specific platform.
 pub type LockedCondaPackages = Vec<RepoDataRecord>;
 
 /// A list of Pypi packages that are locked for a specific platform.
 pub type LockedPypiPackages = Vec<PypiRecord>;
 
-/// A single Pypi record that contains both the package data and the environment data. In Pixi we
-/// basically always need both.
+/// A single Pypi record that contains both the package data and the environment
+/// data. In Pixi we basically always need both.
 pub type PypiRecord = (PypiPackageData, PypiPackageEnvironmentData);
 
-/// Loads the lockfile for the specified project or returns a dummy one if none could be found.
+/// Loads the lockfile for the specified project or returns a dummy one if none
+/// could be found.
 pub async fn load_lock_file(project: &Project) -> miette::Result<LockFile> {
     let lock_file_path = project.lock_file_path();
     if lock_file_path.is_file() {

src/lock_file/update.rs

@@ -93,11 +93,11 @@ pub struct LockFileDerivedData<'p> {
     pub package_cache: PackageCache,
 
     /// A list of prefixes that are up-to-date with the latest conda packages.
-    pub updated_conda_prefixes: HashMap<Environment<'p>, (Prefix, PythonStatus)>,
+    pub updated_conda_prefixes: HashMap<EnvironmentName, (Prefix, PythonStatus)>,
 
     /// A list of prefixes that have been updated while resolving all
     /// dependencies.
-    pub updated_pypi_prefixes: HashMap<Environment<'p>, Prefix>,
+    pub updated_pypi_prefixes: HashMap<EnvironmentName, Prefix>,
 
     /// The cached uv context
     pub uv_context: Option<UvResolutionContext>,
@@ -125,7 +125,7 @@ impl<'p> LockFileDerivedData<'p> {
             },
         )?;
 
-        if let Some(prefix) = self.updated_pypi_prefixes.get(environment) {
+        if let Some(prefix) = self.updated_pypi_prefixes.get(environment.name()) {
             return Ok(prefix.clone());
         }
 
@@ -176,7 +176,7 @@ impl<'p> LockFileDerivedData<'p> {
 
         // Store that we updated the environment, so we won't have to do it again.
         self.updated_pypi_prefixes
-            .insert(environment.clone(), prefix.clone());
+            .insert(environment.name().clone(), prefix.clone());
 
         Ok(prefix)
     }
@@ -210,7 +210,7 @@ impl<'p> LockFileDerivedData<'p> {
         environment: &Environment<'p>,
     ) -> miette::Result<(Prefix, PythonStatus)> {
         // If we previously updated this environment, early out.
-        if let Some((prefix, python_status)) = self.updated_conda_prefixes.get(environment) {
+        if let Some((prefix, python_status)) = self.updated_conda_prefixes.get(environment.name()) {
             return Ok((prefix.clone(), python_status.clone()));
         }
 
@@ -257,8 +257,10 @@ impl<'p> LockFileDerivedData<'p> {
         .await?;
 
         // Store that we updated the environment, so we won't have to do it again.
-        self.updated_conda_prefixes
-            .insert(environment.clone(), (prefix.clone(), python_status.clone()));
+        self.updated_conda_prefixes.insert(
+            environment.name().clone(),
+            (prefix.clone(), python_status.clone()),
+        );
 
         Ok((prefix, python_status))
     }
@@ -434,7 +436,7 @@ impl<'p> UpdateContext<'p> {
     /// Get a list of conda prefixes that have been updated.
     pub fn take_instantiated_conda_prefixes(
         &mut self,
-    ) -> HashMap<Environment<'p>, (Prefix, PythonStatus)> {
+    ) -> HashMap<EnvironmentName, (Prefix, PythonStatus)> {
         self.instantiated_conda_prefixes
             .drain()
             .filter_map(|(env, cell)| match env {
@@ -443,7 +445,7 @@ impl<'p> UpdateContext<'p> {
                         .expect("prefixes must not be shared")
                         .into_inner()
                         .expect("prefix must be available");
-                    Some((env, prefix))
+                    Some((env.name().clone(), prefix))
                 }
                 _ => None,
             })

src/lock_file/utils.rs

@@ -0,0 +1,45 @@
+use rattler_conda_types::Platform;
+use rattler_lock::{LockFile, LockFileBuilder, Package};
+
+use crate::{
+    project::{grouped_environment::GroupedEnvironment, Environment},
+    HasFeatures, Project,
+};
+
+/// Constructs a new lock-file where some of the packages have been removed
+pub fn filter_lock_file<'p, F: FnMut(&Environment<'p>, Platform, &Package) -> bool>(
+    project: &'p Project,
+    lock_file: &LockFile,
+    mut filter: F,
+) -> LockFile {
+    let mut builder = LockFileBuilder::new();
+
+    for (environment_name, environment) in lock_file.environments() {
+        // Find the environment in the project
+        let Some(project_env) = project.environment(environment_name) else {
+            continue;
+        };
+
+        // Copy the channels
+        builder.set_channels(environment_name, environment.channels().to_vec());
+
+        // Copy the indexes
+        let indexes = environment.pypi_indexes().cloned().unwrap_or_else(|| {
+            GroupedEnvironment::from(project_env.clone())
+                .pypi_options()
+                .into()
+        });
+        builder.set_pypi_indexes(environment_name, indexes);
+
+        // Copy all packages that don't need to be relaxed
+        for (platform, packages) in environment.packages_by_platform() {
+            for package in packages {
+                if filter(&project_env, platform, &package) {
+                    builder.add_package(environment_name, platform, package);
+                }
+            }
+        }
+    }
+
+    builder.finish()
+}

src/project/manifest/document.rs

@@ -251,6 +251,7 @@ impl ManifestSource {
         requirement: &pep508_rs::Requirement,
         platform: Option<Platform>,
         feature_name: &FeatureName,
+        editable: Option<bool>,
     ) -> Result<(), TomlError> {
         match self {
             ManifestSource::PyProjectToml(_) => {
@@ -270,10 +271,15 @@ impl ManifestSource {
                 }
             }
             ManifestSource::PixiToml(_) => {
+                let mut pypi_requirement = PyPiRequirement::from(requirement.clone());
+                if let Some(editable) = editable {
+                    pypi_requirement.set_editable(editable);
+                }
+
                 self.get_or_insert_toml_table(platform, feature_name, consts::PYPI_DEPENDENCIES)?
                     .insert(
                         requirement.name.as_ref(),
-                        Item::Value(PyPiRequirement::from(requirement.clone()).into()),
+                        Item::Value(pypi_requirement.into()),
                     );
             }
         };