refactor parsing of packages arguments

src/cli/global/common.rs

@@ -3,7 +3,8 @@ use std::path::PathBuf;
 use indexmap::IndexMap;
 use miette::IntoDiagnostic;
 use rattler_conda_types::{
-    Channel, ChannelConfig, MatchSpec, PackageName, Platform, PrefixRecord, RepoDataRecord,
+    Channel, ChannelConfig, MatchSpec, PackageName, ParseStrictness, Platform, PrefixRecord,
+    RepoDataRecord,
 };
 use rattler_repodata_gateway::sparse::SparseRepoData;
 use rattler_solve::{resolvo, ChannelPriority, SolverImpl, SolverTask};
@@ -16,6 +17,22 @@ use crate::{
     utils::reqwest::build_reqwest_clients,
 };
 
+/// A trait to facilitate extraction of packages data from arguments
+pub(super) trait HasSpecs {
+    /// returns packages passed as arguments to the command
+    fn packages(&self) -> Vec<&str>;
+
+    fn specs(&self) -> miette::Result<IndexMap<PackageName, MatchSpec>> {
+        let mut map = IndexMap::with_capacity(self.packages().len());
+        for package in self.packages() {
+            let spec = MatchSpec::from_str(package, ParseStrictness::Strict).into_diagnostic()?;
+            let name = package_name(&spec)?;
+            map.insert(name, spec);
+        }
+
+        Ok(map)
+    }
+}
 /// Global binaries directory, default to `$HOME/.pixi/bin`
 pub struct BinDir(pub PathBuf);
 

src/cli/global/install.rs

@@ -10,9 +10,7 @@ use itertools::Itertools;
 use miette::IntoDiagnostic;
 use rattler::install::Transaction;
 use rattler::package_cache::PackageCache;
-use rattler_conda_types::{
-    MatchSpec, PackageName, ParseStrictness, Platform, PrefixRecord, RepoDataRecord,
-};
+use rattler_conda_types::{PackageName, Platform, PrefixRecord, RepoDataRecord};
 use rattler_shell::{
     activation::{ActivationVariables, Activator, PathModificationBehavior},
     shell::Shell,
@@ -22,7 +20,7 @@ use reqwest_middleware::ClientWithMiddleware;
 
 use super::common::{
     channel_name_from_prefix, find_designated_package, get_client_and_sparse_repodata,
-    load_package_records, package_name, BinDir, BinEnvDir,
+    load_package_records, BinDir, BinEnvDir, HasSpecs,
 };
 
 /// Installs the defined package in a global accessible location.
@@ -48,6 +46,12 @@ pub struct Args {
     config: ConfigCli,
 }
 
+impl HasSpecs for Args {
+    fn packages(&self) -> Vec<&str> {
+        self.package.iter().map(AsRef::as_ref).collect()
+    }
+}
+
 /// Create the environment activation script
 fn create_activation_script(prefix: &Prefix, shell: ShellEnum) -> miette::Result<String> {
     let activator =
@@ -244,22 +248,13 @@ pub async fn execute(args: Args) -> miette::Result<()> {
     let config = Config::with_cli_config(&args.config);
     let channels = config.compute_channels(&args.channel).into_diagnostic()?;
 
-    // Find the MatchSpec we want to install
-    let specs = args
-        .package
-        .into_iter()
-        .map(|package_str| MatchSpec::from_str(&package_str, ParseStrictness::Strict))
-        .collect::<Result<Vec<_>, _>>()
-        .into_diagnostic()?;
-
     // Fetch sparse repodata
     let (authenticated_client, sparse_repodata) =
         get_client_and_sparse_repodata(&channels, &config).await?;
 
     // Install the package(s)
     let mut executables = vec![];
-    for package_matchspec in specs {
-        let package_name = package_name(&package_matchspec)?;
+    for (package_name, package_matchspec) in args.specs()? {
         let records = load_package_records(package_matchspec, &sparse_repodata)?;
 
         let (prefix_package, scripts, _) =

src/cli/global/remove.rs

@@ -4,12 +4,11 @@ use clap::Parser;
 use clap_verbosity_flag::{Level, Verbosity};
 use itertools::Itertools;
 use miette::IntoDiagnostic;
-use rattler_conda_types::ParseStrictness::Strict;
-use rattler_conda_types::{MatchSpec, PackageName};
+use rattler_conda_types::PackageName;
 
 use crate::prefix::Prefix;
 
-use super::common::{find_designated_package, BinDir, BinEnvDir};
+use super::common::{find_designated_package, BinDir, BinEnvDir, HasSpecs};
 use super::install::{find_and_map_executable_scripts, BinScriptMapping};
 
 /// Removes a package previously installed into a globally accessible location via `pixi global install`.
@@ -24,24 +23,14 @@ pub struct Args {
     verbose: Verbosity,
 }
 
+impl HasSpecs for Args {
+    fn packages(&self) -> Vec<&str> {
+        self.package.iter().map(AsRef::as_ref).collect()
+    }
+}
+
 pub async fn execute(args: Args) -> miette::Result<()> {
-    // Find the MatchSpec we want to remove
-    let specs = args
-        .package
-        .into_iter()
-        .map(|package_str| MatchSpec::from_str(&package_str, Strict))
-        .collect::<Result<Vec<_>, _>>()
-        .into_diagnostic()?;
-    let packages = specs
-        .into_iter()
-        .map(|spec| {
-            spec.name
-                .clone()
-                .ok_or_else(|| miette::miette!("could not find package name in MatchSpec {}", spec))
-        })
-        .collect::<Result<Vec<_>, _>>()?;
-
-    for package_name in packages {
+    for (package_name, _) in args.specs()? {
         remove_global_package(package_name, &args.verbose).await?;
     }
 

src/cli/global/upgrade.rs

@@ -2,18 +2,18 @@ use std::collections::{HashMap, HashSet};
 use std::time::Duration;
 
 use clap::Parser;
+use indexmap::IndexMap;
 use indicatif::ProgressBar;
 use itertools::Itertools;
 use miette::IntoDiagnostic;
-use rattler_conda_types::{Channel, MatchSpec, PackageName, Version};
-use rattler_conda_types::{ParseStrictness, RepoDataRecord};
+use rattler_conda_types::{Channel, MatchSpec, PackageName, RepoDataRecord, Version};
 use reqwest_middleware::ClientWithMiddleware;
 
 use crate::config::Config;
 use crate::progress::{global_multi_progress, long_running_progress_style};
 
 use super::common::{
-    find_installed_package, get_client_and_sparse_repodata, load_package_records, package_name,
+    find_installed_package, get_client_and_sparse_repodata, load_package_records, HasSpecs,
 };
 use super::install::globally_install_package;
 use super::list::list_global_packages;
@@ -39,26 +39,24 @@ pub struct Args {
     channel: Vec<String>,
 }
 
+impl HasSpecs for Args {
+    fn packages(&self) -> Vec<&str> {
+        self.specs.iter().map(AsRef::as_ref).collect()
+    }
+}
+
 pub async fn execute(args: Args) -> miette::Result<()> {
     let config = Config::load_global();
 
     // Get the MatchSpec(s) we need to upgrade
-    let specs = args
-        .specs
-        .iter()
-        .map(|p| MatchSpec::from_str(p, ParseStrictness::Strict).into_diagnostic())
-        .collect::<Result<Vec<_>, _>>()?;
-    let names = specs
-        .iter()
-        .map(package_name)
-        .collect::<Result<Vec<_>, _>>()?;
+    let specs = args.specs()?;
 
     // Return with error if any package is not globally installed.
     let global_packages = list_global_packages()
         .await?
         .into_iter()
         .collect::<HashSet<_>>();
-    let requested = names.iter().cloned().collect::<HashSet<_>>();
+    let requested = specs.keys().cloned().collect::<HashSet<_>>();
     let not_installed = requested.difference(&global_packages).collect_vec();
     match not_installed.len() {
         0 => {} // Do nothing when all packages are globally installed
@@ -72,7 +70,7 @@ pub async fn execute(args: Args) -> miette::Result<()> {
         ),
     };
 
-    upgrade_packages(names, specs, config, &args.channel).await
+    upgrade_packages(specs, config, &args.channel).await
 }
 
 pub(super) async fn upgrade_package(
@@ -103,16 +101,15 @@ pub(super) async fn upgrade_package(
 }
 
 pub(super) async fn upgrade_packages(
-    names: Vec<PackageName>,
-    specs: Vec<MatchSpec>,
+    specs: IndexMap<PackageName, MatchSpec>,
     config: Config,
     cli_channels: &[String],
 ) -> Result<(), miette::Error> {
     // Get channels and versions of globally installed packages
-    let mut installed_versions = HashMap::with_capacity(names.len());
+    let mut installed_versions = HashMap::with_capacity(specs.len());
     let mut channels = config.compute_channels(cli_channels).into_diagnostic()?;
 
-    for package_name in names.iter() {
+    for package_name in specs.keys() {
         let prefix_record = find_installed_package(package_name).await?;
         let last_installed_channel = Channel::from_str(
             prefix_record.repodata_record.channel.clone(),
@@ -137,7 +134,7 @@ pub(super) async fn upgrade_packages(
 
     // Upgrade each package when relevant
     let mut upgraded = false;
-    for (package_name, package_matchspec) in names.into_iter().zip(specs.into_iter()) {
+    for (package_name, package_matchspec) in specs {
         let matchspec_has_version = package_matchspec.version.is_some();
         let records = load_package_records(package_matchspec, &sparse_repodata)?;
         let package_record = records

src/cli/global/upgrade_all.rs

@@ -1,5 +1,5 @@
 use clap::Parser;
-use itertools::Itertools;
+use indexmap::IndexMap;
 
 use rattler_conda_types::MatchSpec;
 
@@ -30,13 +30,16 @@ pub async fn execute(args: Args) -> miette::Result<()> {
     let config = Config::with_cli_config(&args.config);
 
     let names = list_global_packages().await?;
-    let specs = names
-        .iter()
-        .map(|name| MatchSpec {
-            name: Some(name.clone()),
-            ..Default::default()
-        })
-        .collect_vec();
-
-    upgrade_packages(names, specs, config, &args.channel).await
+    let mut specs = IndexMap::with_capacity(names.len());
+    for name in names {
+        specs.insert(
+            name.clone(),
+            MatchSpec {
+                name: Some(name),
+                ..Default::default()
+            },
+        );
+    }
+
+    upgrade_packages(specs, config, &args.channel).await
 }