feat: add --priority arg to project channel add (#2086)

Co-authored-by: Ruben Arts <ruben.arts@hotmail.com>

crates/pixi_manifest/src/channel.rs

@@ -3,6 +3,7 @@ use std::str::FromStr;
 use rattler_conda_types::NamedChannelOrUrl;
 use serde::{de::Error, Deserialize, Deserializer};
 use serde_with::serde_as;
+use toml_edit::{Table, Value};
 
 /// A channel with an optional priority.
 /// If the priority is not specified, it is assumed to be 0.
@@ -23,6 +24,29 @@ impl From<NamedChannelOrUrl> for PrioritizedChannel {
     }
 }
 
+impl From<(NamedChannelOrUrl, Option<i32>)> for PrioritizedChannel {
+    fn from((value, prio): (NamedChannelOrUrl, Option<i32>)) -> Self {
+        Self {
+            channel: value,
+            priority: prio,
+        }
+    }
+}
+
+impl From<PrioritizedChannel> for Value {
+    fn from(channel: PrioritizedChannel) -> Self {
+        match channel.priority {
+            Some(priority) => {
+                let mut table = Table::new().into_inline_table();
+                table.insert("channel", channel.channel.to_string().into());
+                table.insert("priority", i64::from(priority).into());
+                Value::InlineTable(table)
+            }
+            None => Value::String(toml_edit::Formatted::new(channel.channel.to_string())),
+        }
+    }
+}
+
 pub enum TomlPrioritizedChannelStrOrMap {
     Map(PrioritizedChannel),
     Str(NamedChannelOrUrl),

crates/pixi_manifest/src/manifests/manifest.rs

@@ -13,7 +13,7 @@ use itertools::Itertools;
 use miette::{miette, IntoDiagnostic, NamedSource, WrapErr};
 use pixi_spec::PixiSpec;
 use rattler_conda_types::{ChannelConfig, MatchSpec, PackageName, Platform, Version};
-use toml_edit::DocumentMut;
+use toml_edit::{DocumentMut, Value};
 
 use crate::{
     consts,
@@ -531,14 +531,25 @@ impl Manifest {
         };
         let to_add: IndexSet<_> = channels.into_iter().collect();
         let new: IndexSet<_> = to_add.difference(current).cloned().collect();
+        let new_channels: IndexSet<_> = new
+            .clone()
+            .into_iter()
+            .map(|channel| channel.channel)
+            .collect();
 
-        // Add the channels to the manifest
+        // clear channels with modified priority
+        current.retain(|c| !new_channels.contains(&c.channel));
+
+        // Add the updated channels to the manifest
         current.extend(new.clone());
+        let current_clone = current.clone();
 
         // Then to the TOML document
         let channels = self.document.get_array_mut("channels", feature_name)?;
-        for channel in new.iter() {
-            channels.push(channel.channel.to_string())
+        // clear and recreate from current list
+        channels.clear();
+        for channel in current_clone.iter() {
+            channels.push(Value::from(channel.clone()));
         }
 
         Ok(())
@@ -555,28 +566,35 @@ impl Manifest {
             FeatureName::Default => &mut self.parsed.project.channels,
             FeatureName::Named(_) => self.feature_mut(feature_name)?.channels_mut(),
         };
-
         // Get the channels to remove, while checking if they exist
         let to_remove: IndexSet<_> = channels
             .into_iter()
             .map(|c| {
                 current
                     .iter()
-                    .position(|x| x.channel == c.channel)
+                    .position(|x| x.channel.to_string() == c.channel.to_string())
                     .ok_or_else(|| miette::miette!("channel {} does not exist", c.channel.as_str()))
-                    .map(|_| c)
+                    .map(|_| c.channel.to_string())
             })
             .collect::<Result<_, _>>()?;
 
-        let retained: IndexSet<_> = current.difference(&to_remove).cloned().collect();
+        let retained: IndexSet<_> = current
+            .iter()
+            .filter(|channel| !to_remove.contains(&channel.channel.to_string()))
+            .cloned()
+            .collect();
 
         // Remove channels from the manifest
         current.retain(|c| retained.contains(c));
+        let current_clone = current.clone();
 
         // And from the TOML document
-        let retained = retained.iter().map(|c| c.channel.as_str()).collect_vec();
         let channels = self.document.get_array_mut("channels", feature_name)?;
-        channels.retain(|x| retained.contains(&x.as_str().unwrap()));
+        // clear and recreate from current list
+        channels.clear();
+        for channel in current_clone.iter() {
+            channels.push(Value::from(channel.clone()));
+        }
 
         Ok(())
     }
@@ -1534,6 +1552,35 @@ platforms = ["linux-64", "win-64"]
             .iter()
             .any(|c| c.channel == custom_channel.channel));
 
+        // Test adding priority
+        let prioritized_channel1 = PrioritizedChannel {
+            channel: NamedChannelOrUrl::Name(String::from("prioritized")),
+            priority: Some(12i32),
+        };
+        manifest
+            .add_channels([prioritized_channel1.clone()], &FeatureName::Default)
+            .unwrap();
+        assert!(manifest
+            .parsed
+            .project
+            .channels
+            .iter()
+            .any(|c| c.channel == prioritized_channel1.channel && c.priority == Some(12i32)));
+
+        let prioritized_channel2 = PrioritizedChannel {
+            channel: NamedChannelOrUrl::Name(String::from("prioritized2")),
+            priority: Some(-12i32),
+        };
+        manifest
+            .add_channels([prioritized_channel2.clone()], &FeatureName::Default)
+            .unwrap();
+        assert!(manifest
+            .parsed
+            .project
+            .channels
+            .iter()
+            .any(|c| c.channel == prioritized_channel2.channel && c.priority == Some(-12i32)));
+
         assert_snapshot!(manifest.document.to_string());
     }
 

crates/pixi_manifest/src/manifests/snapshots/pixi_manifest__manifests__manifest__tests__add_channels.snap

@@ -4,7 +4,7 @@ expression: manifest.document.to_string()
 ---
 [project]
 name = "foo"
-channels = ["conda-forge", "https://custom.com/channel"]
+channels = ["conda-forge", "https://custom.com/channel", { channel = "prioritized", priority = 12 }, { channel = "prioritized2", priority = -12 }]
 platforms = ["linux-64", "win-64"]
 
 [dependencies]

src/cli/project/channel/mod.rs

@@ -26,6 +26,10 @@ pub struct AddRemoveArgs {
     #[clap(required = true, num_args=1..)]
     pub channel: Vec<NamedChannelOrUrl>,
 
+    /// Specify the channel priority
+    #[clap(long, num_args = 1)]
+    pub priority: Option<i32>,
+
     /// Don't update the environment, only modify the manifest and the
     /// lock-file.
     #[clap(long)]
@@ -38,7 +42,10 @@ pub struct AddRemoveArgs {
 
 impl AddRemoveArgs {
     fn prioritized_channels(&self) -> impl IntoIterator<Item = PrioritizedChannel> + '_ {
-        self.channel.iter().cloned().map(PrioritizedChannel::from)
+        self.channel
+            .iter()
+            .cloned()
+            .map(|channel| PrioritizedChannel::from((channel, self.priority)))
     }
 
     fn feature_name(&self) -> FeatureName {
@@ -51,15 +58,19 @@ impl AddRemoveArgs {
         for channel in self.channel {
             match channel {
                 NamedChannelOrUrl::Name(ref name) => eprintln!(
-                    "{}{operation} {} ({})",
+                    "{}{operation} {} ({}){}",
                     console::style(console::Emoji("✔ ", "")).green(),
                     name,
-                    channel.clone().into_base_url(channel_config)
+                    channel.clone().into_base_url(channel_config),
+                    self.priority
+                        .map_or_else(|| "".to_string(), |p| format!(" at priority {}", p))
                 ),
                 NamedChannelOrUrl::Url(url) => eprintln!(
-                    "{}{operation} {}",
+                    "{}{operation} {}{}",
                     console::style(console::Emoji("✔ ", "")).green(),
-                    url
+                    url,
+                    self.priority
+                        .map_or_else(|| "".to_string(), |p| format!(" at priority {}", p)),
                 ),
             }
         }

tests/common/builders.rs

@@ -323,6 +323,11 @@ impl ProjectChannelAddBuilder {
         self
     }
 
+    pub fn with_priority(mut self, priority: Option<i32>) -> Self {
+        self.args.priority = priority;
+        self
+    }
+
     /// Alias to add a local channel.
     pub fn with_local_channel(self, channel: impl AsRef<Path>) -> Self {
         self.with_channel(Url::from_directory_path(channel).unwrap())
@@ -342,6 +347,39 @@ impl IntoFuture for ProjectChannelAddBuilder {
     }
 }
 
+pub struct ProjectChannelRemoveBuilder {
+    pub manifest_path: Option<PathBuf>,
+    pub args: project::channel::AddRemoveArgs,
+}
+
+impl ProjectChannelRemoveBuilder {
+    /// Removes the specified channel
+    pub fn with_channel(mut self, name: impl Into<String>) -> Self {
+        self.args
+            .channel
+            .push(NamedChannelOrUrl::from_str(&name.into()).unwrap());
+        self
+    }
+
+    /// Alias to Remove a local channel.
+    pub fn with_local_channel(self, channel: impl AsRef<Path>) -> Self {
+        self.with_channel(Url::from_directory_path(channel).unwrap())
+    }
+}
+
+impl IntoFuture for ProjectChannelRemoveBuilder {
+    type Output = miette::Result<()>;
+    type IntoFuture = Pin<Box<dyn Future<Output = Self::Output> + 'static>>;
+
+    fn into_future(self) -> Self::IntoFuture {
+        project::channel::execute(project::channel::Args {
+            manifest_path: self.manifest_path,
+            command: project::channel::Command::Remove(self.args),
+        })
+        .boxed_local()
+    }
+}
+
 /// Contains the arguments to pass to [`install::execute()`]. Call `.await` to
 /// call the CLI execute method and await the result at the same time.
 pub struct InstallBuilder {

tests/common/mod.rs

@@ -38,7 +38,7 @@ use thiserror::Error;
 
 use self::builders::{HasDependencyConfig, RemoveBuilder};
 use crate::common::builders::{
-    AddBuilder, InitBuilder, InstallBuilder, ProjectChannelAddBuilder,
+    AddBuilder, InitBuilder, InstallBuilder, ProjectChannelAddBuilder, ProjectChannelRemoveBuilder,
     ProjectEnvironmentAddBuilder, TaskAddBuilder, TaskAliasBuilder, UpdateBuilder,
 };
 
@@ -334,6 +334,20 @@ impl PixiControl {
                 channel: vec![],
                 no_install: true,
                 feature: None,
+                priority: None,
+            },
+        }
+    }
+
+    /// Add a new channel to the project.
+    pub fn project_channel_remove(&self) -> ProjectChannelRemoveBuilder {
+        ProjectChannelRemoveBuilder {
+            manifest_path: Some(self.manifest_path()),
+            args: project::channel::AddRemoveArgs {
+                channel: vec![],
+                no_install: true,
+                feature: None,
+                priority: None,
             },
         }
     }

tests/project_tests.rs

@@ -12,7 +12,7 @@ use url::Url;
 use crate::common::{package_database::PackageDatabase, PixiControl};
 
 #[tokio::test]
-async fn add_channel() {
+async fn add_remove_channel() {
     // Create a local package database with no entries and write it to disk. This
     // ensures that we have a valid channel.
     let package_database = PackageDatabase::default();
@@ -47,7 +47,40 @@ async fn add_channel() {
     let local_channel =
         NamedChannelOrUrl::Url(Url::from_file_path(additional_channel_dir.as_ref()).unwrap());
     let channels = project.default_environment().channels();
+    assert!(channels.len() == 2);
+    assert!(channels.last().unwrap() == &&local_channel);
     assert!(channels.contains(&local_channel));
+
+    // now add the same channel, with priority 2
+    pixi.project_channel_add()
+        .with_local_channel(additional_channel_dir.path())
+        .with_priority(Some(2i32))
+        .await
+        .unwrap();
+
+    // Load again
+    let project = Project::from_path(&pixi.manifest_path()).unwrap();
+    let channels = project.default_environment().channels();
+    // still present
+    assert!(channels.contains(&local_channel));
+    // didn't duplicate
+    assert!(channels.len() == 2);
+    // priority applied
+    assert!(channels.first().unwrap() == &&local_channel);
+
+    // now remove it
+    pixi.project_channel_remove()
+        .with_local_channel(additional_channel_dir.path())
+        .await
+        .unwrap();
+
+    // Load again
+    let project = Project::from_path(&pixi.manifest_path()).unwrap();
+    let channels = project.default_environment().channels();
+
+    // Channel has been removed
+    assert!(channels.len() == 1);
+    assert!(!channels.contains(&local_channel));
 }
 
 #[tokio::test]