feat: Accept sparse repodata in solve

crates/rattler_repodata_gateway/src/gateway/local_subdir.rs

@@ -22,6 +22,12 @@ pub struct LocalSubdirClient {
 }
 
 impl LocalSubdirClient {
+    /// Create a client directly from an already-loaded [`SparseRepoData`],
+    /// without parsing anything from disk.
+    pub fn new(sparse: Arc<SparseRepoData>) -> Self {
+        Self { sparse }
+    }
+
     pub fn from_file(
         repodata_path: &Path,
         channel: Channel,

crates/rattler_repodata_gateway/src/gateway/mod.rs

@@ -1876,6 +1876,79 @@ mod test {
         );
     }
 
+    /// Loads the `dummy` test channel's `linux-64` subdir as a `SparseRepoData`.
+    /// `foobar` depends on `bors`, which is used to exercise recursive queries.
+    fn dummy_sparse_repo_data() -> crate::sparse::SparseRepoData {
+        let channel_config = ChannelConfig::default_with_root_dir(std::env::current_dir().unwrap());
+        let channel = Channel::from_str("dummy", &channel_config).unwrap();
+        let path = Path::new(env!("CARGO_MANIFEST_DIR"))
+            .join("../../test-data/channels/dummy/linux-64/repodata.json");
+        crate::sparse::SparseRepoData::from_file(channel, "linux-64", path, None).unwrap()
+    }
+
+    #[tokio::test]
+    async fn test_sparse_repodata_source() {
+        let gateway = Gateway::new();
+        let source = Arc::new(dummy_sparse_repo_data());
+
+        // Query the sparse repodata directly for `foobar`.
+        let records = gateway
+            .query(
+                vec![super::Source::SparseRepoData(vec![source.clone()])],
+                vec![Platform::Linux64],
+                vec![PackageName::from_str("foobar").unwrap()].into_iter(),
+            )
+            .recursive(false)
+            .await
+            .unwrap();
+
+        let all_records: Vec<_> = records.iter().flat_map(RepoData::iter).collect();
+        assert!(!all_records.is_empty(), "should have foobar records");
+        assert!(
+            all_records
+                .iter()
+                .all(|r| r.package_record.name.as_normalized() == "foobar"),
+            "non-recursive query should only return foobar records"
+        );
+
+        // A recursive query should also pull in `bors`, `foobar`'s dependency.
+        let records = gateway
+            .query(
+                vec![super::Source::SparseRepoData(vec![source.clone()])],
+                vec![Platform::Linux64],
+                vec![PackageName::from_str("foobar").unwrap()].into_iter(),
+            )
+            .recursive(true)
+            .await
+            .unwrap();
+
+        let all_records: Vec<_> = records.iter().flat_map(RepoData::iter).collect();
+        assert!(
+            all_records
+                .iter()
+                .any(|r| r.package_record.name.as_normalized() == "bors"),
+            "recursive query should also fetch foobar's dependency bors"
+        );
+
+        // A `SparseRepoData` only ever represents the one channel/subdir pair it was
+        // loaded from; other platforms should yield no records.
+        let records = gateway
+            .query(
+                vec![super::Source::SparseRepoData(vec![source])],
+                vec![Platform::Win64],
+                vec![PackageName::from_str("foobar").unwrap()].into_iter(),
+            )
+            .recursive(false)
+            .await
+            .unwrap();
+
+        let all_records: Vec<_> = records.iter().flat_map(RepoData::iter).collect();
+        assert!(
+            all_records.is_empty(),
+            "querying a platform other than the sparse repodata's own subdir should be empty"
+        );
+    }
+
     #[tokio::test]
     async fn test_mixed_channel_and_custom_source() {
         let gateway = Gateway::new();

crates/rattler_repodata_gateway/src/gateway/query.rs

@@ -15,6 +15,7 @@ use super::{
     BarrierCell, ChannelNoticeResult, GatewayError, GatewayInner, GatewayWarning, RepoData,
     channel_expander::{ChannelExpander, ChannelRelationsMode, ChannelRelationsWarning},
     channel_relations::DEFAULT_CHANNEL_RELATIONS_MAX_DEPTH,
+    local_subdir::LocalSubdirClient,
     source::{CustomSourceClient, Source},
     subdir::{PackageRecords, Subdir, SubdirData},
 };
@@ -536,6 +537,40 @@ impl QueryExecutor {
                         });
                         (SubdirKind::Custom, fut)
                     }
+                    Source::SparseRepoData(sparse_list) => {
+                        // Each entry represents a different subdir, so find the one
+                        // matching the requested platform; if none matches, treat it
+                        // as having no records, same as a channel that doesn't
+                        // publish a given subdir.
+                        let matching = sparse_list
+                            .iter()
+                            .find(|sparse| platform.as_str() == sparse.subdir())
+                            .cloned();
+                        let url = matching
+                            .as_ref()
+                            .or_else(|| sparse_list.first())
+                            .map(|sparse| sparse.channel.base_url.clone());
+                        let kind = match url {
+                            Some(url) => SubdirKind::Channel { url, platform },
+                            None => SubdirKind::Custom,
+                        };
+                        let subdir = match matching {
+                            Some(sparse) => Arc::new(Subdir::Found(SubdirData::from_client(
+                                LocalSubdirClient::new(sparse),
+                            ))),
+                            None => Arc::new(Subdir::NotFound),
+                        };
+                        let b = barrier.clone();
+                        let fut = box_future(async move {
+                            b.set(subdir.clone()).expect("subdir was set twice");
+                            Ok(PendingSubdirOk {
+                                subdir,
+                                kind_url_and_platform: None,
+                                warning: None,
+                            })
+                        });
+                        (kind, fut)
+                    }
                 };
 
                 subdir_handles.push(SubdirHandle {

crates/rattler_repodata_gateway/src/gateway/source.rs

@@ -8,7 +8,7 @@ use super::{
     GatewayError,
     subdir::{PackageRecords, SubdirClient, extract_unique_deps_split},
 };
-use crate::Reporter;
+use crate::{Reporter, sparse::SparseRepoData};
 
 /// A source of repodata records for a specific subdirectory.
 ///
@@ -39,14 +39,18 @@ pub trait RepoDataSource: Send + Sync {
 /// A source of repodata, either a channel or a custom source.
 ///
 /// This enum allows the [`Gateway::query()`](super::Gateway::query) method
-/// to accept both traditional channels and custom repodata sources.
+/// to accept both traditional channels custom repodata sources and sparse repodata.
 #[derive(Clone)]
 pub enum Source {
     /// A traditional conda channel (expanded to all requested platforms).
     Channel(Channel),
 
     /// A custom repodata source (provides records for requested platforms).
     Custom(Arc<dyn RepoDataSource>),
+
+    /// A sparse repodata source (provides records for requested platforms from sparse
+    /// repodata). Each entry represents a different subdir.
+    SparseRepoData(Vec<Arc<SparseRepoData>>),
 }
 
 impl From<Channel> for Source {
@@ -61,6 +65,18 @@ impl From<Arc<dyn RepoDataSource>> for Source {
     }
 }
 
+impl From<Arc<SparseRepoData>> for Source {
+    fn from(source: Arc<SparseRepoData>) -> Self {
+        Source::SparseRepoData(vec![source])
+    }
+}
+
+impl From<Vec<Arc<SparseRepoData>>> for Source {
+    fn from(sources: Vec<Arc<SparseRepoData>>) -> Self {
+        Source::SparseRepoData(sources)
+    }
+}
+
 /// Adapts a [`RepoDataSource`] to the internal [`SubdirClient`] trait
 /// for a specific platform.
 ///

crates/rattler_repodata_gateway/src/sparse/mod.rs

@@ -82,7 +82,7 @@ pub struct SparseRepoData {
     inner: SparseRepoDataInner,
 
     /// The channel from which this data was downloaded.
-    channel: Channel,
+    pub channel: Channel,
 
     /// The subdirectory from where the repodata is downloaded
     subdir: String,

py-rattler/rattler/repo_data/gateway.py

@@ -469,11 +469,14 @@ def _convert_sources(sources: Iterable[Any]) -> List[Any]:
     Channels are converted to their internal PyChannel representation.
     Custom RepoDataSource implementations are wrapped in an adapter that
     converts between FFI types and Python wrapper types.
+    SparseRepoData objects are converted to their internal PySparseRepoData
+    representation.
 
     Raises:
         TypeError: If a source doesn't implement the required interface.
     """
     from rattler.repo_data.source import RepoDataSource
+    from rattler.repo_data.sparse import SparseRepoData
 
     converted = []
     for source in sources:
@@ -483,12 +486,15 @@ def _convert_sources(sources: Iterable[Any]) -> List[Any]:
         elif isinstance(source, Channel):
             # Channel object - extract PyChannel
             converted.append(source._channel)
+        elif isinstance(source, SparseRepoData):
+            # SparseRepoData object - extract PySparseRepoData
+            converted.append(source._sparse)
         elif isinstance(source, RepoDataSource):
             # Wrap RepoDataSource in adapter for FFI type conversion
             converted.append(_RepoDataSourceAdapter(source))
         else:
             raise TypeError(
-                f"Expected Channel, str, or object implementing RepoDataSource protocol, "
+                f"Expected Channel, str, SparseRepoData, or object implementing RepoDataSource protocol, "
                 f"got {type(source).__name__}. "
                 f"See rattler.RepoDataSource for the required interface."
             )

py-rattler/rattler/solver/solver.py

@@ -22,7 +22,7 @@
 
 
 async def solve(
-    sources: Sequence[Union[Channel, str, RepoDataSource]],
+    sources: Sequence[Union[Channel, str, RepoDataSource, SparseRepoData]],
     specs: Sequence[MatchSpec | str],
     gateway: Gateway = Gateway(),
     platforms: Optional[Sequence[Platform | PlatformLiteral]] = None,
@@ -43,7 +43,7 @@ async def solve(
 
     Arguments:
         sources: The sources to query for the packages. Can be channels (by name, URL,
-                 or Channel object) or custom RepoDataSource implementations.
+                 or Channel object), custom RepoDataSource implementations or SparseRepoData objects.
         specs: A list of matchspec to solve.
         platforms: The platforms to query for the packages. If `None` the current platform and
                 `noarch` is used.

py-rattler/src/repo_data/gateway.rs

@@ -5,7 +5,9 @@ use std::sync::Arc;
 use pyo3::exceptions::{PyTypeError, PyValueError};
 use pyo3::pybacked::PyBackedStr;
 use pyo3::types::PyAnyMethods;
-use pyo3::{Borrowed, Bound, FromPyObject, PyAny, PyErr, PyResult, Python, pyclass, pymethods};
+use pyo3::{
+    Borrowed, Bound, FromPyObject, PyAny, PyErr, PyRef, PyResult, Python, pyclass, pymethods,
+};
 use pyo3_async_runtimes::tokio::future_into_py;
 use rattler_repodata_gateway::fetch::{CacheAction, FetchRepoDataOptions, Variant};
 use rattler_repodata_gateway::{
@@ -22,6 +24,7 @@ use crate::platform::PyPlatform;
 use crate::record::PyRecord;
 use crate::repo_data::PyChannelRelations;
 use crate::repo_data::source::PyRepoDataSource;
+use crate::repo_data::sparse::PySparseRepoData;
 use crate::{PyChannel, Wrap};
 
 #[pyclass(from_py_object)]
@@ -146,6 +149,7 @@ pub(crate) fn emit_gateway_warnings(warnings: Vec<GatewayWarning>) -> PyResult<(
 ///
 /// Accepts either:
 /// - A `PyChannel` object (wrapped Channel)
+/// - A `PySparseRepoData` object
 /// - Any object implementing the `RepoDataSource` protocol
 ///   (has `fetch_package_records` and `package_names` methods)
 pub fn py_object_to_source(obj: Bound<'_, PyAny>) -> PyResult<Source> {
@@ -154,6 +158,11 @@ pub fn py_object_to_source(obj: Bound<'_, PyAny>) -> PyResult<Source> {
         return Ok(Source::from(channel.inner));
     }
 
+    // Then try to extract as SparseRepoData
+    if let Ok(sparse) = obj.extract::<PyRef<'_, PySparseRepoData>>() {
+        return Ok(Source::from(sparse.as_source()?));
+    }
+
     // Check if it implements the RepoDataSource protocol
     if obj.hasattr("fetch_package_records")? && obj.hasattr("package_names")? {
         let source = PyRepoDataSource::new(obj.unbind());
@@ -163,7 +172,7 @@ pub fn py_object_to_source(obj: Bound<'_, PyAny>) -> PyResult<Source> {
     }
 
     Err(PyTypeError::new_err(
-        "Expected Channel or object implementing RepoDataSource protocol \
+        "Expected Channel, SparseRepoData, or object implementing RepoDataSource protocol \
          (with fetch_package_records and package_names methods)",
     ))
 }
@@ -360,11 +369,14 @@ impl PyGateway {
         // Separate channels and custom sources
         let mut channels: Vec<rattler_conda_types::Channel> = Vec::new();
         let mut custom_sources: Vec<Arc<dyn rattler_repodata_gateway::RepoDataSource>> = Vec::new();
+        let mut sparse_sources: Vec<Arc<rattler_repodata_gateway::sparse::SparseRepoData>> =
+            Vec::new();
 
         for source in rust_sources {
             match source {
                 Source::Channel(channel) => channels.push(channel),
                 Source::Custom(custom) => custom_sources.push(custom),
+                Source::SparseRepoData(sparse) => sparse_sources.extend(sparse),
             }
         }
 

py-rattler/src/repo_data/sparse.rs

@@ -22,24 +22,38 @@ pub struct PySparseRepoData {
     // in a RwLock because most of the time we just want to be able to read from it. We only
     // need write access to close it.
     //
-    // This whole thing is then wrapped in an Arc so we can share this with a background thread
-    // without blocking the GIL.
-    pub(crate) inner: Arc<RwLock<Option<SparseRepoData>>>,
-    subdir: String,
+    // The `SparseRepoData` itself is wrapped in an `Arc` too, so `as_source` can hand out a
+    // cheap clone of the *same* `Arc` on every call rather than re-parsing or copying data.
+    //
+    // This whole thing is then wrapped in an outer Arc so we can share this with a background
+    // thread without blocking the GIL.
+    pub(crate) inner: Arc<RwLock<Option<Arc<SparseRepoData>>>>,
+    pub(crate) subdir: String,
 }
 
 impl PySparseRepoData {
     /// Create a new instance without requiring the GIL.
     pub(crate) fn from_args(channel: PyChannel, subdir: String, path: PathBuf) -> PyResult<Self> {
         Ok(SparseRepoData::from_file(channel.into(), subdir, path, None)?.into())
     }
+
+    /// Returns the underlying `SparseRepoData` so it can be passed directly
+    /// to `Gateway::query` as a `Source::SparseRepoData` (e.g. via `solve`'s
+    /// `sources` argument).
+    pub(crate) fn as_source(&self) -> PyResult<Arc<SparseRepoData>> {
+        self.inner
+            .read()
+            .as_ref()
+            .cloned()
+            .ok_or_else(|| PyValueError::new_err("I/O operation on closed file."))
+    }
 }
 
 impl From<SparseRepoData> for PySparseRepoData {
     fn from(value: SparseRepoData) -> Self {
         Self {
             subdir: value.subdir().to_owned(),
-            inner: Arc::new(RwLock::new(Some(value))),
+            inner: Arc::new(RwLock::new(Some(Arc::new(value)))),
         }
     }
 }
@@ -228,7 +242,7 @@ impl PySparseRepoData {
         let repo_data_refs = repo_data_locks
             .iter()
             .map(|s| {
-                s.as_ref()
+                s.as_deref()
                     .ok_or_else(|| PyValueError::new_err("I/O operation on closed file."))
             })
             .collect::<Result<Vec<_>, _>>()?;

py-rattler/src/solver.rs

@@ -200,7 +200,7 @@ pub fn py_solve_with_sparse_repodata<'py>(
             let repo_data_refs = repo_data_locks
                 .iter()
                 .map(|s| {
-                    s.as_ref()
+                    s.as_deref()
                         .ok_or_else(|| PyValueError::new_err("I/O operation on closed file."))
                 })
                 .collect::<Result<Vec<_>, _>>()?;