feat(repodata_gateway): read virtual package plugin registrations

crates/rattler-bin/Cargo.toml

@@ -34,6 +34,10 @@ rustls = [
 s3 = ["rattler_networking/s3", "rattler_upload/s3"]
 gcs = ["rattler_networking/gcs"]
 oauth = ["rattler/oauth"]
+experimental-virtual-package-plugins = [
+  "rattler_repodata_gateway/experimental-virtual-package-plugins",
+  "rattler_index/experimental-virtual-package-plugins",
+]
 
 [dependencies]
 anyhow = { workspace = true }

crates/rattler-bin/src/commands/virtual_packages.rs

@@ -1,17 +1,111 @@
 use miette::IntoDiagnostic;
 use rattler_conda_types::GenericVirtualPackage;
+#[cfg(feature = "experimental-virtual-package-plugins")]
+use rattler_conda_types::Platform;
 use rattler_virtual_packages::VirtualPackageOverrides;
 
 /// Print detected virtual packages.
 #[derive(Debug, clap::Parser)]
-pub struct Opt {}
+#[cfg_attr(
+    feature = "experimental-virtual-package-plugins",
+    clap(after_help = r#"Examples:
+  rattler virtual-packages
+  rattler virtual-packages -c ./test-data/channels/virtual-package-plugins"#)
+)]
+pub struct Opt {
+    /// Channels to list registered virtual package plugins for
+    #[cfg(feature = "experimental-virtual-package-plugins")]
+    #[clap(short, long)]
+    channels: Vec<String>,
 
-pub fn virtual_packages(_opt: Opt) -> miette::Result<()> {
+    /// Platforms to read registrations for [default: current and noarch]
+    #[cfg(feature = "experimental-virtual-package-plugins")]
+    #[clap(short, long)]
+    platforms: Vec<Platform>,
+}
+
+pub async fn virtual_packages(opt: Opt, offline: bool) -> miette::Result<()> {
     let virtual_packages =
         rattler_virtual_packages::VirtualPackage::detect(&VirtualPackageOverrides::from_env())
             .into_diagnostic()?;
     for package in virtual_packages {
         println!("{}", GenericVirtualPackage::from(package.clone()));
     }
+
+    #[cfg(feature = "experimental-virtual-package-plugins")]
+    print_plugins(&opt.channels, &opt.platforms, offline).await?;
+
+    #[cfg(not(feature = "experimental-virtual-package-plugins"))]
+    let _ = (opt, offline);
+
+    Ok(())
+}
+
+/// Prints the plugin registrations declared by each `(channel, platform)`
+/// subdirectory, in the order the channels were given.
+#[cfg(feature = "experimental-virtual-package-plugins")]
+async fn print_plugins(
+    channels: &[String],
+    platforms: &[Platform],
+    offline: bool,
+) -> miette::Result<()> {
+    use std::{collections::HashMap, env};
+
+    use itertools::Itertools;
+    use rattler_conda_types::{Channel, ChannelConfig, PackageName};
+    use rattler_repodata_gateway::{Gateway, SourceConfig};
+
+    if channels.is_empty() {
+        return Ok(());
+    }
+
+    let channel_config =
+        ChannelConfig::default_with_root_dir(env::current_dir().into_diagnostic()?);
+    let channels = channels
+        .iter()
+        .map(|channel| Channel::from_str(channel, &channel_config))
+        .collect::<Result<Vec<_>, _>>()
+        .into_diagnostic()?;
+
+    let platforms = if platforms.is_empty() {
+        vec![Platform::current(), Platform::NoArch]
+    } else {
+        platforms.to_vec()
+    };
+
+    let gateway = Gateway::builder()
+        .with_client(super::client::create_client_with_middleware(offline)?)
+        .with_channel_config(rattler_repodata_gateway::ChannelConfig {
+            default: SourceConfig {
+                cache_action: super::client::repodata_cache_action(offline),
+                ..SourceConfig::default()
+            },
+            per_channel: HashMap::new(),
+        })
+        .finish();
+
+    for channel in &channels {
+        for platform in &platforms {
+            let plugins = gateway
+                .virtual_package_plugins(channel, *platform)
+                .await
+                .into_diagnostic()?;
+            if plugins.is_empty() {
+                continue;
+            }
+            println!(
+                "\nvirtual package plugins in {} [{platform}]:",
+                channel.canonical_name()
+            );
+            for (plugin, provided) in &plugins {
+                println!(
+                    "  {} provides {}",
+                    plugin.as_source(),
+                    provided.iter().map(PackageName::as_source).join(", ")
+                );
+            }
+        }
+    }
+
     Ok(())
 }

crates/rattler-bin/src/main.rs

@@ -124,7 +124,9 @@ async fn async_main() -> miette::Result<()> {
         Command::Solve(opts) => commands::solve::solve(opts, offline).await,
         Command::List(opts) => commands::list::list(opts).await,
         Command::ShellHook(opts) => commands::shell_hook::shell_hook(opts).await,
-        Command::VirtualPackages(opts) => commands::virtual_packages::virtual_packages(opts),
+        Command::VirtualPackages(opts) => {
+            commands::virtual_packages::virtual_packages(opts, offline).await
+        }
         Command::InstallMenu(opts) => commands::menu::install_menu(opts).await,
         Command::RemoveMenu(opts) => commands::menu::remove_menu(opts).await,
         Command::Run(opts) => commands::run::run(opts).await,

crates/rattler_conda_types/Cargo.toml

@@ -13,6 +13,7 @@ readme.workspace = true
 
 [features]
 default = ["rayon"]
+experimental-virtual-package-plugins = []
 
 [package.metadata.docs.rs]
 all-features = true

crates/rattler_conda_types/src/lib.rs

@@ -64,6 +64,8 @@ pub use platform::{Arch, ParseArchError, ParsePlatformError, Platform};
 pub use prefix_data::PrefixData;
 pub use prefix_record::PrefixRecord;
 pub use record_traits::HasArtifactIdentificationRefs;
+#[cfg(feature = "experimental-virtual-package-plugins")]
+pub use repo_data::VirtualPackagePlugins;
 pub use repo_data::{
     ChannelInfo, ChannelRelations, ConvertSubdirError, PackageRecord, RecordFromPath, RepoData,
     RepodataRevision, RepodataRevisionInfo, RepodataRevisionMetadata, RepodataRevisions,

crates/rattler_conda_types/src/repo_data/mod.rs

@@ -107,8 +107,26 @@ pub struct ChannelInfo {
     /// [CEP-42](https://github.com/conda/ceps/blob/main/cep-0042.md).
     #[serde(default, skip_serializing_if = "ChannelRelations::is_none_or_empty")]
     pub channel_relations: Option<ChannelRelations>,
+
+    /// Virtual package detection plugins registered by the channel.
+    #[cfg(feature = "experimental-virtual-package-plugins")]
+    #[serde_as(
+        deserialize_as = "IndexMap<DeserializeFromStrUnchecked, Vec<DeserializeFromStrUnchecked>>"
+    )]
+    #[serde(default, skip_serializing_if = "IndexMap::is_empty")]
+    pub virtual_package_plugins: VirtualPackagePlugins,
 }
 
+/// Virtual package detection plugins registered by a channel: the name of the
+/// package providing the plugin, mapped to the virtual packages it provides.
+///
+/// One plugin may provide several virtual packages, e.g. a `cuda-detect`
+/// providing both `__cuda` and `__cuda_arch`. The executable to run is named
+/// after the plugin package. Inverting the map is left to the caller: the
+/// reverse direction is many-to-many.
+#[cfg(feature = "experimental-virtual-package-plugins")]
+pub type VirtualPackagePlugins = IndexMap<PackageName, Vec<PackageName>>;
+
 /// Repodata revisions keyed by revision, mirroring the `vN` dictionary of the
 /// CEP draft <https://github.com/conda/ceps/pull/146>. Keying encodes
 /// uniqueness; insertion order is preserved.
@@ -1084,6 +1102,8 @@ mod test {
         package::DistArchiveIdentifier,
         repo_data::{compute_package_url, determine_subdir},
     };
+    #[cfg(feature = "experimental-virtual-package-plugins")]
+    use crate::{PackageName, repo_data::VirtualPackagePlugins};
 
     // isl-0.12.2-1.tar.bz2
     // gmp-5.1.2-6.tar.bz2
@@ -1166,6 +1186,8 @@ mod test {
                     base: Some("../conda-forge".to_string()),
                     overrides: None,
                 }),
+                #[cfg(feature = "experimental-virtual-package-plugins")]
+                virtual_package_plugins: VirtualPackagePlugins::default(),
             }),
             packages: IndexMap::default(),
             conda_packages: IndexMap::default(),
@@ -1188,6 +1210,8 @@ mod test {
                     base_url: None,
                     repodata_revisions: IndexMap::default(),
                     channel_relations,
+                    #[cfg(feature = "experimental-virtual-package-plugins")]
+                    virtual_package_plugins: VirtualPackagePlugins::default(),
                 }),
                 packages: IndexMap::default(),
                 conda_packages: IndexMap::default(),
@@ -1199,6 +1223,99 @@ mod test {
         }
     }
 
+    #[cfg(feature = "experimental-virtual-package-plugins")]
+    #[test]
+    fn test_virtual_package_plugins() {
+        // A single plugin may provide several virtual packages; the order of
+        // both the plugins and their virtual packages is preserved.
+        let raw = r#"{
+            "info": {
+                "subdir": "linux-64",
+                "virtual_package_plugins": {
+                    "cuda-detect": ["__cuda", "__cuda_arch"],
+                    "rocm-detect": ["__rocm"]
+                }
+            },
+            "packages": {},
+            "packages.conda": {}
+        }"#;
+        let repodata: RepoData = serde_json::from_str(raw).unwrap();
+        let plugins = &repodata.info.as_ref().unwrap().virtual_package_plugins;
+
+        assert_eq!(
+            plugins
+                .keys()
+                .map(PackageName::as_source)
+                .collect::<Vec<_>>(),
+            ["cuda-detect", "rocm-detect"]
+        );
+        assert_eq!(
+            plugins[&PackageName::new_unchecked("cuda-detect")]
+                .iter()
+                .map(PackageName::as_source)
+                .collect::<Vec<_>>(),
+            ["__cuda", "__cuda_arch"]
+        );
+        assert_eq!(
+            plugins[&PackageName::new_unchecked("rocm-detect")]
+                .iter()
+                .map(PackageName::as_source)
+                .collect::<Vec<_>>(),
+            ["__rocm"]
+        );
+
+        let json = serde_json::to_string(&repodata).unwrap();
+        assert!(json.contains("\"virtual_package_plugins\""));
+        assert_eq!(serde_json::from_str::<RepoData>(&json).unwrap(), repodata);
+
+        let without = RepoData {
+            version: Some(2),
+            info: Some(ChannelInfo {
+                subdir: Some("linux-64".to_string()),
+                base_url: None,
+                repodata_revisions: IndexMap::default(),
+                channel_relations: None,
+                virtual_package_plugins: VirtualPackagePlugins::default(),
+            }),
+            packages: IndexMap::default(),
+            conda_packages: IndexMap::default(),
+            v3: V3Packages::default(),
+            removed: ahash::HashSet::default(),
+        };
+        assert!(
+            !serde_json::to_string(&without)
+                .unwrap()
+                .contains("virtual_package_plugins")
+        );
+    }
+
+    /// Names are deserialized unchecked, so a channel publishing a malformed
+    /// name does not render the entire repodata unusable.
+    #[cfg(feature = "experimental-virtual-package-plugins")]
+    #[test]
+    fn test_virtual_package_plugins_malformed_name() {
+        let raw = r#"{
+            "info": {
+                "subdir": "linux-64",
+                "virtual_package_plugins": { "invalid$plugin": ["__rocm"] }
+            },
+            "packages": {},
+            "packages.conda": {}
+        }"#;
+        let repodata: RepoData = serde_json::from_str(raw).unwrap();
+        assert_eq!(
+            repodata
+                .info
+                .as_ref()
+                .unwrap()
+                .virtual_package_plugins
+                .keys()
+                .map(PackageName::as_source)
+                .collect::<Vec<_>>(),
+            ["invalid$plugin"]
+        );
+    }
+
     #[test]
     fn test_repodata_revisions() {
         let raw = r#"{

crates/rattler_conda_types/src/repo_data/sharded.rs

@@ -2,7 +2,11 @@
 
 use crate::PackageRecord;
 use crate::package::DistArchiveIdentifier;
+#[cfg(feature = "experimental-virtual-package-plugins")]
+use crate::repo_data::VirtualPackagePlugins;
 use crate::repo_data::{ChannelRelations, RepodataRevisions, V3Packages};
+#[cfg(feature = "experimental-virtual-package-plugins")]
+use crate::utils::serde::DeserializeFromStrUnchecked;
 use crate::utils::serde::{sort_index_map_alphabetically, sort_set_alphabetically};
 use indexmap::IndexMap;
 use jiff::Timestamp;
@@ -60,6 +64,14 @@ pub struct ShardedSubdirInfo {
     /// [CEP-42](https://github.com/conda/ceps/blob/main/cep-0042.md).
     #[serde(default, skip_serializing_if = "ChannelRelations::is_none_or_empty")]
     pub channel_relations: Option<ChannelRelations>,
+
+    /// Virtual package detection plugins registered by the channel.
+    #[cfg(feature = "experimental-virtual-package-plugins")]
+    #[serde_as(
+        deserialize_as = "IndexMap<DeserializeFromStrUnchecked, Vec<DeserializeFromStrUnchecked>>"
+    )]
+    #[serde(default, skip_serializing_if = "IndexMap::is_empty")]
+    pub virtual_package_plugins: VirtualPackagePlugins,
 }
 
 #[cfg(test)]
@@ -131,11 +143,48 @@ mod tests {
                 created_at: None,
                 repodata_revisions: IndexMap::default(),
                 channel_relations,
+                #[cfg(feature = "experimental-virtual-package-plugins")]
+                virtual_package_plugins: VirtualPackagePlugins::default(),
             };
             let json = serde_json::to_string(&info).unwrap();
             assert!(!json.contains("channel_relations"));
         }
     }
+
+    #[cfg(feature = "experimental-virtual-package-plugins")]
+    #[test]
+    fn test_sharded_subdir_info_virtual_package_plugins() {
+        let raw = r#"{
+            "subdir": "linux-64",
+            "base_url": "./",
+            "shards_base_url": "./shards/",
+            "virtual_package_plugins": {
+                "cuda-detect": ["__cuda", "__cuda_arch"]
+            }
+        }"#;
+        let info: ShardedSubdirInfo = serde_json::from_str(raw).unwrap();
+        assert_eq!(
+            info.virtual_package_plugins[&PackageName::new_unchecked("cuda-detect")]
+                .iter()
+                .map(PackageName::as_source)
+                .collect::<Vec<_>>(),
+            ["__cuda", "__cuda_arch"]
+        );
+
+        let json = serde_json::to_string(&info).unwrap();
+        assert!(json.contains("\"virtual_package_plugins\""));
+
+        // Omitted entirely when no plugins are registered.
+        let info = ShardedSubdirInfo {
+            virtual_package_plugins: VirtualPackagePlugins::default(),
+            ..info
+        };
+        assert!(
+            !serde_json::to_string(&info)
+                .unwrap()
+                .contains("virtual_package_plugins")
+        );
+    }
 }
 
 /// An individual shard that contains repodata for a single package name.

crates/rattler_index/Cargo.toml

@@ -28,6 +28,10 @@ rustls = [
   "opendal/reqwest-rustls-tls",
 ]
 s3 = ["opendal/services-s3", "dep:rattler_s3"]
+experimental-virtual-package-plugins = [
+  "rattler_conda_types/experimental-virtual-package-plugins",
+  "rattler_repodata_gateway/experimental-virtual-package-plugins",
+]
 
 [[bin]]
 name = "rattler-index"