Add --pypi option to pixi remove

src/cli/remove.rs

@@ -1,17 +1,19 @@
 use std::path::PathBuf;
+use std::str::FromStr;
 
 use clap::Parser;
-use rattler_conda_types::{PackageName, Platform};
+use miette::miette;
+use rattler_conda_types::Platform;
 
 use crate::environment::LockFileUsage;
-use crate::{environment::get_up_to_date_prefix, project::SpecType, Project};
+use crate::{consts, environment::get_up_to_date_prefix, project::SpecType, Project};
 
 /// Remove the dependency from the project
 #[derive(Debug, Default, Parser)]
 pub struct Args {
     /// List of dependencies you wish to remove from the project
     #[arg(required = true)]
-    pub deps: Vec<PackageName>,
+    pub deps: Vec<String>,
 
     /// The path to 'pixi.toml'
     #[arg(long)]
@@ -25,11 +27,27 @@ pub struct Args {
     #[arg(long, conflicts_with = "host")]
     pub build: bool,
 
+    /// Whether the dependency is a pypi package
+    #[arg(long)]
+    pub pypi: bool,
+
     /// The platform for which the dependency should be removed
     #[arg(long, short)]
     pub platform: Option<Platform>,
 }
 
+fn convert_pkg_name<T>(deps: &[String]) -> miette::Result<Vec<T>>
+where
+    T: FromStr,
+{
+    deps.iter()
+        .map(|dep| {
+            T::from_str(dep)
+                .map_err(|_| miette!("Can't convert dependency name `{dep}` to package name"))
+        })
+        .collect()
+}
+
 pub async fn execute(args: Args) -> miette::Result<()> {
     let mut project = Project::load_or_else_discover(args.manifest_path.as_deref())?;
     let deps = args.deps;
@@ -41,43 +59,60 @@ pub async fn execute(args: Args) -> miette::Result<()> {
         SpecType::Run
     };
 
-    let results = deps
-        .iter()
-        .map(|dep| {
-            if let Some(p) = &args.platform {
+    let section_name: String = if args.pypi {
+        consts::PYPI_DEPENDENCIES.to_string()
+    } else {
+        spec_type.name().to_string()
+    };
+    let table_name = if let Some(p) = &args.platform {
+        format!("target.{}.{}", p.as_str(), section_name)
+    } else {
+        section_name
+    };
+
+    fn format_ok_message(pkg_name: &String, pkg_extras: &String, table_name: &String) -> String {
+        format!(
+            "Removed {} from [{}]",
+            console::style(format!("{pkg_name} {pkg_extras}")).bold(),
+            console::style(table_name).bold()
+        )
+    }
+    let mut sucessful_output: Vec<String> = Vec::with_capacity(deps.len());
+    if args.pypi {
+        let all_pkg_name = convert_pkg_name::<rip::types::PackageName>(&deps)?;
+        for dep in all_pkg_name.iter() {
+            let result = if let Some(p) = &args.platform {
+                project.manifest.remove_target_pypi_dependency(dep, p)?
+            } else {
+                project.manifest.remove_pypi_dependency(dep)?
+            };
+            sucessful_output.push(format_ok_message(
+                &result.0.as_str().to_string(),
+                &result.1.to_string(),
+                &table_name,
+            ));
+        }
+    } else {
+        let all_pkg_name = convert_pkg_name::<rattler_conda_types::PackageName>(&deps)?;
+        for dep in all_pkg_name.iter() {
+            let result = if let Some(p) = &args.platform {
                 project
                     .manifest
-                    .remove_target_dependency(dep, &spec_type, p)
+                    .remove_target_dependency(dep, &spec_type, p)?
             } else {
-                project.manifest.remove_dependency(dep, &spec_type)
-            }
-        })
-        .collect::<Vec<_>>();
-
+                project.manifest.remove_dependency(dep, &spec_type)?
+            };
+            sucessful_output.push(format_ok_message(
+                &result.0,
+                &result.1.to_string(),
+                &table_name,
+            ));
+        }
+    };
     project.save()?;
+    eprintln!("{}", sucessful_output.join("\n"));
 
     // updating prefix after removing from toml
     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 {
-            format!("target.{}.{}", p.as_str(), spec_type.name())
-        } else {
-            spec_type.name().to_string()
-        };
-
-        eprintln!(
-            "Removed {} from [{}]",
-            console::style(format!("{removed} {spec}")).bold(),
-            console::style(table_name).bold(),
-        );
-    }
-
-    for result in &results {
-        if let Err(e) = result {
-            eprintln!("{e}");
-        }
-    }
-
     Ok(())
 }

src/consts.rs

@@ -3,3 +3,4 @@ pub const PROJECT_LOCK_FILE: &str = "pixi.lock";
 pub const PIXI_DIR: &str = ".pixi";
 pub const PREFIX_FILE_NAME: &str = "prefix";
 pub const ENVIRONMENT_DIR: &str = "env";
+pub const PYPI_DEPENDENCIES: &str = "pypi-dependencies";

src/project/manifest.rs

@@ -468,6 +468,23 @@ impl Manifest {
         Ok(())
     }
 
+    pub fn remove_pypi_dependency(
+        &mut self,
+        dep: &rip::types::PackageName,
+    ) -> miette::Result<(rip::types::PackageName, PyPiRequirement)> {
+        if let Item::Table(ref mut t) = self.document[consts::PYPI_DEPENDENCIES] {
+            if t.contains_key(dep.as_str()) && t.remove(dep.as_str()).is_some() {
+                return self.parsed.remove_pypi_dependencies(dep.as_str());
+            }
+        }
+
+        Err(miette::miette!(
+            "Couldn't find {} in {}",
+            console::style(dep.as_str()).bold(),
+            consts::PYPI_DEPENDENCIES,
+        ))
+    }
+
     /// Removes a dependency from `pixi.toml` based on `SpecType`.
     pub fn remove_dependency(
         &mut self,
@@ -502,6 +519,18 @@ impl Manifest {
             .remove_target_dependency(dep.as_normalized(), spec_type, platform)
     }
 
+    /// Removes a target specific dependency from `pixi.toml` based on pypi-dependencies
+    pub fn remove_target_pypi_dependency(
+        &mut self,
+        dep: &rip::types::PackageName,
+        platform: &Platform,
+    ) -> miette::Result<(rip::types::PackageName, PyPiRequirement)> {
+        let table = get_toml_target_table(&mut self.document, platform, consts::PYPI_DEPENDENCIES)?;
+        table.remove(dep.as_str());
+        self.parsed
+            .remove_target_pypi_dependency(dep.as_str(), platform)
+    }
+
     /// Returns a mutable reference to the channels array.
     fn channels_array_mut(&mut self) -> miette::Result<&mut Array> {
         let project = &mut self.document["project"];
@@ -802,6 +831,29 @@ impl ProjectManifest {
         }
     }
 
+    /// Remove PyPi dependencies
+    pub fn remove_pypi_dependencies(
+        &mut self,
+        dep: &str,
+    ) -> miette::Result<(PackageName, PyPiRequirement)> {
+        self.pypi_dependencies
+            .as_mut()
+            .and_then(|pypi_deps| {
+                let key_to_remove = pypi_deps
+                    .keys()
+                    .find(|&pkg_name| pkg_name.as_str() == dep)
+                    .cloned();
+
+                key_to_remove.and_then(|pkg_key| pypi_deps.shift_remove_entry(&pkg_key))
+            })
+            .ok_or_else(|| {
+                miette::miette!(
+                    "Couldn't remove PyPi dependency: {}",
+                    console::style(dep).bold(),
+                )
+            })
+    }
+
     /// Remove dependency given a `SpecType`.
     pub fn remove_dependency(
         &mut self,
@@ -828,6 +880,57 @@ impl ProjectManifest {
         }
     }
 
+    pub fn remove_target_pypi_dependency(
+        &mut self,
+        dep: &str,
+        platform: &Platform,
+    ) -> miette::Result<(PackageName, PyPiRequirement)> {
+        let target = PixiSpanned::from(TargetSelector::Platform(*platform));
+        let target_metadata = self.target.get_mut(&target).ok_or(miette::miette!(
+            "Platform: {} is not configured for this project",
+            console::style(platform.as_str()).bold(),
+        ))?;
+        let dependencies = target_metadata.pypi_dependencies.as_mut();
+
+        if let Some(deps) = dependencies {
+            let key_to_remove = deps.keys().find(|k| k.as_str() == dep);
+            if let Some(key_dep) = key_to_remove {
+                deps.shift_remove_entry(&key_dep.clone())
+                    .ok_or(miette::miette!(
+                        "Couldn't find {} in [{}]",
+                        console::style(dep).bold(),
+                        console::style(format!(
+                            "target.{}.{}",
+                            platform.as_str(),
+                            consts::PYPI_DEPENDENCIES
+                        ))
+                        .bold(),
+                    ))
+            } else {
+                Err(miette::miette!(
+                    "Couldn't find {} in [{}]",
+                    console::style(dep).bold(),
+                    console::style(format!(
+                        "target.{}.{}",
+                        platform.as_str(),
+                        consts::PYPI_DEPENDENCIES
+                    ))
+                    .bold(),
+                ))
+            }
+        } else {
+            Err(miette::miette!(
+                "[{}] doesn't exist",
+                console::style(format!(
+                    "target.{}.{}",
+                    platform.as_str(),
+                    consts::PYPI_DEPENDENCIES
+                ))
+                .bold(),
+            ))
+        }
+    }
+
     /// Remove a dependency for a `Platform`.
     pub fn remove_target_dependency(
         &mut self,
@@ -1444,6 +1547,73 @@ mod test {
         assert_debug_snapshot!(manifest.parsed.target);
     }
 
+    #[test]
+    fn test_remove_target_pypi_dependency() {
+        let pixi_cfg = r#"[project]
+name = "pixi_fun"
+version = "0.1.0"
+channels = []
+platforms = ["linux-64", "win-64"]
+
+[dependencies]
+python = ">=3.12.1,<3.13"
+
+[target.win-64.pypi-dependencies]
+jax = { version = "*", extras = ["cpu"] }
+requests = "*"
+
+[target.linux-64.pypi-dependencies]
+xpackage = "==1.2.3"
+ypackage = {version = ">=1.2.3"}
+"#;
+        let tmpdir = tempdir().unwrap();
+        let mut manifest = Manifest::from_str(tmpdir.path(), pixi_cfg).unwrap();
+        manifest
+            .parsed
+            .remove_target_pypi_dependency("xpackage", &Platform::Linux64)
+            .unwrap();
+        manifest
+            .parsed
+            .remove_target_pypi_dependency("jax", &Platform::Win64)
+            .unwrap();
+        assert_debug_snapshot!(manifest.parsed);
+    }
+
+    #[test]
+    fn test_remove_pypi_dependencies() {
+        let pixi_cfg = r#"[project]
+name = "pixi_fun"
+version = "0.1.0"
+channels = []
+platforms = ["linux-64", "win-64"]
+
+[dependencies]
+python = ">=3.12.1,<3.13"
+
+[pypi-dependencies]
+jax = { version = "*", extras = ["cpu"] }
+requests = "*"
+xpackage = "==1.2.3"
+ypackage = {version = ">=1.2.3"}
+"#;
+        let tmpdir = tempdir().unwrap();
+        let mut manifest = Manifest::from_str(tmpdir.path(), pixi_cfg).unwrap();
+        manifest
+            .remove_pypi_dependency(&rip::types::PackageName::from_str("jax").unwrap())
+            .unwrap();
+        let unwrapped_pypi = manifest.parsed.pypi_dependencies.as_ref().unwrap();
+        assert!(!unwrapped_pypi.contains_key(&rip::types::PackageName::from_str("jax").unwrap()));
+        assert!(
+            unwrapped_pypi.contains_key(&rip::types::PackageName::from_str("requests").unwrap())
+        );
+        assert!(
+            unwrapped_pypi.contains_key(&rip::types::PackageName::from_str("xpackage").unwrap())
+        );
+        assert!(
+            unwrapped_pypi.contains_key(&rip::types::PackageName::from_str("ypackage").unwrap())
+        );
+    }
+
     #[test]
     fn test_set_version() {
         // Using known files in the project so the test succeed including the file check.

src/project/mod.rs

@@ -41,7 +41,7 @@ impl DependencyType {
     pub fn name(&self) -> &'static str {
         match self {
             DependencyType::CondaDependency(dep) => dep.name(),
-            DependencyType::PypiDependency => "pypi-dependencies",
+            DependencyType::PypiDependency => consts::PYPI_DEPENDENCIES,
         }
     }
 }