feat: Allow to upgrade several global packages at once (#1324)

This PR refactors most of `upgrade.rs` and `upgrade_all.rs` into a
common `upgrade_packages` function, which upgrades several global
packages at once.

Now, several global packages can be upgraded at once.

---------

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

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);
 
@@ -108,7 +125,7 @@ pub fn bin_env_dir() -> Option<PathBuf> {
 /// # Returns
 ///
 /// The package name from the given MatchSpec
-pub(super) fn package_name(package_matchspec: &MatchSpec) -> miette::Result<PackageName> {
+fn package_name(package_matchspec: &MatchSpec) -> miette::Result<PackageName> {
     package_matchspec.name.clone().ok_or_else(|| {
         miette::miette!(
             "could not find package name in MatchSpec {}",

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/list.rs

@@ -148,11 +148,9 @@ pub(super) async fn list_global_packages() -> miette::Result<Vec<PackageName>> {
 
     while let Some(entry) = dir_contents.next_entry().await.into_diagnostic()? {
         if entry.file_type().await.into_diagnostic()?.is_dir() {
-            let Ok(name) = PackageName::from_str(entry.file_name().to_string_lossy().as_ref())
-            else {
-                continue;
-            };
-            packages.push(name);
+            if let Ok(name) = PackageName::from_str(entry.file_name().to_string_lossy().as_ref()) {
+                packages.push(name);
+            }
         }
     }
 

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

@@ -1,28 +1,28 @@
+use std::collections::HashMap;
 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 reqwest_middleware::ClientWithMiddleware;
+use rattler_conda_types::{Channel, MatchSpec, PackageName};
 
 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;
 
 /// Upgrade specific package which is installed globally.
 #[derive(Parser, Debug)]
 #[clap(arg_required_else_help = true)]
 pub struct Args {
-    /// Specifies the package that is to be upgraded.
-    package: String,
+    /// Specifies the packages to upgrade.
+    #[arg(required = true)]
+    pub specs: Vec<String>,
 
     /// Represents the channels from which to upgrade specified package.
     /// Multiple channels can be specified by using this field multiple times.
@@ -37,111 +37,97 @@ pub struct Args {
     channel: Vec<String>,
 }
 
-pub async fn execute(args: Args) -> miette::Result<()> {
-    // Get the MatchSpec we need to upgrade
-    let package_matchspec =
-        MatchSpec::from_str(&args.package, ParseStrictness::Strict).into_diagnostic()?;
-    let package_name = package_name(&package_matchspec)?;
-    let matchspec_has_version = package_matchspec.version.is_some();
-
-    // Return with error if this package is not globally installed.
-    if !list_global_packages()
-        .await?
-        .iter()
-        .any(|global_package| global_package.as_normalized() == package_name.as_normalized())
-    {
-        miette::bail!(
-            "Package {} is not globally installed",
-            package_name.as_source()
-        );
+impl HasSpecs for Args {
+    fn packages(&self) -> Vec<&str> {
+        self.specs.iter().map(AsRef::as_ref).collect()
     }
+}
 
-    let prefix_record = find_installed_package(&package_name).await?;
-    let installed_version = prefix_record
-        .repodata_record
-        .package_record
-        .version
-        .into_version();
-
+pub async fn execute(args: Args) -> miette::Result<()> {
     let config = Config::load_global();
+    let specs = args.specs()?;
+    upgrade_packages(specs, config, &args.channel).await
+}
 
-    // Figure out what channels we are using
-    let last_installed_channel = Channel::from_str(
-        prefix_record.repodata_record.channel.clone(),
-        config.channel_config(),
-    )
-    .into_diagnostic()?;
-
-    let mut channels = vec![last_installed_channel];
-    let input_channels = args
-        .channel
-        .iter()
-        .map(|c| Channel::from_str(c, config.channel_config()))
-        .collect::<Result<Vec<Channel>, _>>()
+pub(super) async fn upgrade_packages(
+    specs: IndexMap<PackageName, MatchSpec>,
+    config: Config,
+    cli_channels: &[String],
+) -> miette::Result<()> {
+    // Get channels and versions of globally installed packages
+    let mut installed_versions = HashMap::with_capacity(specs.len());
+    let mut channels = config.compute_channels(cli_channels).into_diagnostic()?;
+
+    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(),
+            config.channel_config(),
+        )
         .into_diagnostic()?;
-    channels.extend(input_channels);
-    // Remove possible duplicates
-    channels = channels.into_iter().unique().collect::<Vec<_>>();
+
+        channels.push(last_installed_channel);
+
+        let installed_version = prefix_record
+            .repodata_record
+            .package_record
+            .version
+            .into_version();
+        installed_versions.insert(package_name.clone(), installed_version);
+    }
+    channels = channels.into_iter().unique().collect();
 
     // Fetch sparse repodata
     let (authenticated_client, sparse_repodata) =
         get_client_and_sparse_repodata(&channels, &config).await?;
 
-    let records = load_package_records(package_matchspec, &sparse_repodata)?;
-    let package_record = records
-        .iter()
-        .find(|r| r.package_record.name.as_normalized() == package_name.as_normalized())
-        .ok_or_else(|| {
-            miette::miette!(
-                "Package {} not found in the specified channels",
-                package_name.as_normalized()
-            )
-        })?;
-    let toinstall_version = package_record.package_record.version.version().to_owned();
-
-    if !matchspec_has_version
-        && toinstall_version.cmp(&installed_version) != std::cmp::Ordering::Greater
-    {
-        eprintln!(
-            "Package {} is already up-to-date",
-            package_name.as_normalized(),
-        );
-        return Ok(());
+    // Upgrade each package when relevant
+    let mut upgraded = false;
+    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
+            .iter()
+            .find(|r| r.package_record.name == package_name)
+            .ok_or_else(|| {
+                miette::miette!(
+                    "Package {} not found in the specified channels",
+                    package_name.as_normalized()
+                )
+            })?;
+        let toinstall_version = package_record.package_record.version.version().to_owned();
+        let installed_version = installed_versions
+            .get(&package_name)
+            .expect("should have the installed version")
+            .to_owned();
+
+        // Perform upgrade if a specific version was requested
+        // OR if a more recent version is available
+        if matchspec_has_version || toinstall_version > installed_version {
+            let message = format!(
+                "{} v{} -> v{}",
+                package_name.as_normalized(),
+                installed_version,
+                toinstall_version
+            );
+
+            let pb = global_multi_progress().add(ProgressBar::new_spinner());
+            pb.enable_steady_tick(Duration::from_millis(100));
+            pb.set_style(long_running_progress_style());
+            pb.set_message(format!(
+                "{} {}",
+                console::style("Updating").green(),
+                message
+            ));
+            globally_install_package(&package_name, records, authenticated_client.clone()).await?;
+            pb.finish_with_message(format!("{} {}", console::style("Updated").green(), message));
+            upgraded = true;
+        }
     }
 
-    upgrade_package(
-        &package_name,
-        installed_version,
-        toinstall_version,
-        records,
-        authenticated_client,
-    )
-    .await
-}
+    if !upgraded {
+        eprintln!("Nothing to upgrade");
+    }
 
-pub(super) async fn upgrade_package(
-    package_name: &PackageName,
-    installed_version: Version,
-    toinstall_version: Version,
-    records: Vec<RepoDataRecord>,
-    authenticated_client: ClientWithMiddleware,
-) -> miette::Result<()> {
-    let message = format!(
-        "{} v{} -> v{}",
-        package_name.as_normalized(),
-        installed_version,
-        toinstall_version
-    );
-
-    let pb = global_multi_progress().add(ProgressBar::new_spinner());
-    pb.enable_steady_tick(Duration::from_millis(100));
-    pb.set_style(long_running_progress_style());
-    pb.set_message(format!(
-        "{} {}",
-        console::style("Updating").green(),
-        message
-    ));
-    globally_install_package(package_name, records, authenticated_client).await?;
-    pb.finish_with_message(format!("{} {}", console::style("Updated").green(), message));
     Ok(())
 }