feat: --pypi for add command

src/cli/add.rs

@@ -4,6 +4,7 @@ use crate::project::SpecType;
 use crate::{
     consts,
     lock_file::{load_lock_file, update_lock_file},
+    project::python::PyPiRequirement,
     project::Project,
     virtual_packages::get_minimal_virtual_packages,
 };
@@ -20,15 +21,17 @@ use rattler_repodata_gateway::sparse::SparseRepoData;
 use rattler_solve::{resolvo, SolverImpl};
 use std::collections::{HashMap, HashSet};
 use std::path::PathBuf;
+use std::str::FromStr;
 
 /// Adds a dependency to the project
 #[derive(Parser, Debug, Default)]
 #[clap(arg_required_else_help = true)]
 pub struct Args {
     /// Specify the dependencies you wish to add to the project.
     ///
-    /// All dependencies should be defined as MatchSpec. If no specific version is
-    /// provided, the latest version compatible with your project will be chosen automatically.
+    /// The dependencies should be defined as MatchSpec, for conda package or a PyPI requirement
+    /// for the --pypi dependencies. If no specific version is provided, the latest version
+    /// compatible with your project will be chosen automatically or a * will be used.
     ///
     /// Example usage:
     ///
@@ -50,8 +53,12 @@ pub struct Args {
     ///
     /// Mixing `--platform` and `--build`/`--host` flags is supported
     ///
+    /// The `--pypi` option will add the package as a pypi-dependency this can not be mixed with the conda dependencies
+    /// - `pixi add --pypi boto3`
+    /// - `pixi add --pypi "boto3==version"
+    ///
     #[arg(required = true)]
-    pub specs: Vec<MatchSpec>,
+    pub specs: Vec<String>,
 
     /// The path to 'pixi.toml'
     #[arg(long)]
@@ -65,6 +72,10 @@ pub struct Args {
     #[arg(long, conflicts_with = "host")]
     pub build: bool,
 
+    /// This is a pypi dependency
+    #[arg(long, conflicts_with_all = ["host", "build"])]
+    pub pypi: bool,
+
     /// Don't update lockfile, implies the no-install as well.
     #[clap(long, conflicts_with = "no_install")]
     pub no_lockfile_update: bool,
@@ -80,7 +91,9 @@ pub struct Args {
 
 impl SpecType {
     pub fn from_args(args: &Args) -> Self {
-        if args.host {
+        if args.pypi {
+            Self::Pypi
+        } else if args.host {
             Self::Host
         } else if args.build {
             Self::Build
@@ -111,18 +124,137 @@ pub async fn execute(args: Args) -> miette::Result<()> {
         .collect::<Vec<Platform>>();
     project.add_platforms(platforms_to_add.iter())?;
 
-    add_specs_to_project(
-        &mut project,
-        args.specs,
-        spec_type,
-        args.no_install,
-        args.no_lockfile_update,
-        spec_platforms,
-    )
-    .await
+    match spec_type {
+        SpecType::Host | SpecType::Build | SpecType::Run => {
+            let specs = args
+                .specs
+                .into_iter()
+                .map(|s| MatchSpec::from_str(&s))
+                .collect::<Result<Vec<_>, _>>()
+                .into_diagnostic()?;
+            add_conda_specs_to_project(
+                &mut project,
+                specs,
+                spec_type,
+                args.no_install,
+                args.no_lockfile_update,
+                spec_platforms,
+            )
+            .await
+        }
+        SpecType::Pypi => {
+            // Parse specs as pep508_rs requirements
+            let pep508_requirements = args
+                .specs
+                .into_iter()
+                .map(|input| pep508_rs::Requirement::from_str(input.as_ref()).into_diagnostic())
+                .collect::<miette::Result<Vec<_>>>()?;
+
+            // Move those requirements into our custom PyPiRequirement
+            let specs_result: miette::Result<Vec<(rip::types::PackageName, PyPiRequirement)>> =
+                pep508_requirements
+                    .into_iter()
+                    .map(|req| {
+                        let name = rip::types::PackageName::from_str(req.name.as_str())?;
+                        let requirement = PyPiRequirement::from(req);
+                        Ok((name, requirement))
+                    })
+                    .collect();
+
+            specs_result
+                .map(|specs| async move {
+                    add_pypi_specs_to_project(
+                        &mut project,
+                        specs,
+                        spec_platforms,
+                        args.no_lockfile_update,
+                        args.no_install,
+                    )
+                    .await
+                })?
+                .await
+        }
+    }
 }
 
-pub async fn add_specs_to_project(
+pub async fn add_pypi_specs_to_project(
+    project: &mut Project,
+    specs: Vec<(rip::types::PackageName, PyPiRequirement)>,
+    specs_platforms: Vec<Platform>,
+    no_update_lockfile: bool,
+    no_install: bool,
+) -> miette::Result<()> {
+    for (name, spec) in &specs {
+        // TODO: Get best version
+        // Add the dependency to the project
+        if specs_platforms.is_empty() {
+            project.add_pypi_dependency(name, spec)?;
+        } else {
+            for platform in specs_platforms.iter() {
+                project.add_target_pypi_dependency(*platform, name.clone(), spec)?;
+            }
+        }
+    }
+    project.save()?;
+
+    let sparse_repo_data = project.fetch_sparse_repodata().await?;
+
+    // Update the lock file
+    let lock_file = if !no_update_lockfile {
+        Some(
+            update_lock_file(
+                project,
+                load_lock_file(project).await?,
+                Some(sparse_repo_data),
+            )
+            .await?,
+        )
+    } else {
+        None
+    };
+
+    if let Some(lock_file) = lock_file {
+        if !no_install {
+            let platform = Platform::current();
+            if project.platforms().contains(&platform) {
+                // Get the currently installed packages
+                let prefix = Prefix::new(project.root().join(".pixi/env"))?;
+                let installed_packages = prefix.find_installed_packages(None).await?;
+
+                // Update the prefix
+                update_prefix(
+                    project.pypi_package_db()?,
+                    &prefix,
+                    installed_packages,
+                    &lock_file,
+                    platform,
+                )
+                .await?;
+            } else {
+                eprintln!("{} skipping installation of environment because your platform ({platform}) is not supported by this project.", style("!").yellow().bold())
+            }
+        }
+    }
+
+    for spec in specs {
+        eprintln!(
+            "{}Added `{}` as pypi dependency",
+            console::style(console::Emoji("✔ ", "")).green(),
+            spec.0.as_str(),
+        );
+    }
+
+    // Print something if we've added for platforms
+    if !specs_platforms.is_empty() {
+        eprintln!(
+            "Added these only for platform(s): {}",
+            specs_platforms.iter().join(", ")
+        )
+    }
+
+    Ok(())
+}
+pub async fn add_conda_specs_to_project(
     project: &mut Project,
     specs: Vec<MatchSpec>,
     spec_type: SpecType,
@@ -158,6 +290,7 @@ pub async fn add_specs_to_project(
             SpecType::Host => project.host_dependencies(platform)?,
             SpecType::Build => project.build_dependencies(platform)?,
             SpecType::Run => project.dependencies(platform)?,
+            _ => return Err(miette::miette!("Fail")),
         };
         // Solve the environment with the new specs added
         let solved_versions = match determine_best_version(
@@ -259,6 +392,7 @@ pub async fn add_specs_to_project(
     match spec_type {
         SpecType::Host => eprintln!("Added these as host dependencies."),
         SpecType::Build => eprintln!("Added these as build dependencies."),
+        SpecType::Pypi => eprintln!("Added these as pypi dependencies."),
         SpecType::Run => {}
     };
 

src/project/manifest.rs

@@ -7,6 +7,7 @@ use indexmap::IndexMap;
 use miette::{Context, IntoDiagnostic, LabeledSpan, NamedSource, Report};
 use rattler_conda_types::{Channel, NamelessMatchSpec, Platform, Version};
 use rattler_virtual_packages::{Archspec, Cuda, LibC, Linux, Osx, VirtualPackage};
+use rip::types::PackageName;
 use serde::Deserializer;
 use serde_with::de::DeserializeAsWrap;
 use serde_with::{serde_as, DeserializeAs, DisplayFromStr, PickFirst};
@@ -143,6 +144,24 @@ impl ProjectManifest {
                     self.build_dependencies.insert(IndexMap::new())
                 }
             }
+            SpecType::Pypi => panic!("Misusing function"),
+        }
+    }
+
+    /// Get the map of dependencies for a given spec type.
+    pub fn create_or_get_pypi_dependencies(
+        &mut self,
+        spec_type: SpecType,
+    ) -> &mut IndexMap<PackageName, PyPiRequirement> {
+        match spec_type {
+            SpecType::Run | SpecType::Host | SpecType::Build => panic!("Misusing function"),
+            SpecType::Pypi => {
+                if let Some(ref mut deps) = self.pypi_dependencies {
+                    deps
+                } else {
+                    self.pypi_dependencies.insert(IndexMap::new())
+                }
+            }
         }
     }
 
@@ -156,6 +175,7 @@ impl ProjectManifest {
             SpecType::Run => Some(&mut self.dependencies),
             SpecType::Build => self.build_dependencies.as_mut(),
             SpecType::Host => self.host_dependencies.as_mut(),
+            _ => None,
         };
 
         if let Some(deps) = dependencies {
@@ -189,6 +209,7 @@ impl ProjectManifest {
             SpecType::Run => Some(&mut target_metadata.dependencies),
             SpecType::Build => target_metadata.build_dependencies.as_mut(),
             SpecType::Host => target_metadata.host_dependencies.as_mut(),
+            _ => None,
         };
 
         if let Some(deps) = dependencies {
@@ -259,6 +280,9 @@ pub struct TargetMetadata {
     #[serde_as(as = "Option<IndexMap<_, PickFirst<(_, DisplayFromStr)>>>")]
     pub build_dependencies: Option<IndexMap<String, NamelessMatchSpec>>,
 
+    #[serde(default, rename = "pypi-dependencies")]
+    pub pypi_dependencies: Option<IndexMap<rip::types::PackageName, PyPiRequirement>>,
+
     /// Additional information to activate an environment.
     #[serde(default)]
     pub activation: Option<Activation>,