diff --git a/crates/rattler_repodata_gateway/src/gateway/local_subdir.rs b/crates/rattler_repodata_gateway/src/gateway/local_subdir.rs index ad1d32bdd0..682c9309b0 100644 --- a/crates/rattler_repodata_gateway/src/gateway/local_subdir.rs +++ b/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) -> Self { + Self { sparse } + } + pub fn from_file( repodata_path: &Path, channel: Channel, diff --git a/crates/rattler_repodata_gateway/src/gateway/mod.rs b/crates/rattler_repodata_gateway/src/gateway/mod.rs index 874068e499..99b02a6808 100644 --- a/crates/rattler_repodata_gateway/src/gateway/mod.rs +++ b/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(); diff --git a/crates/rattler_repodata_gateway/src/gateway/query.rs b/crates/rattler_repodata_gateway/src/gateway/query.rs index e6344c54b5..be4f64aade 100644 --- a/crates/rattler_repodata_gateway/src/gateway/query.rs +++ b/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 { diff --git a/crates/rattler_repodata_gateway/src/gateway/source.rs b/crates/rattler_repodata_gateway/src/gateway/source.rs index ce1f94d428..ff68a53e4f 100644 --- a/crates/rattler_repodata_gateway/src/gateway/source.rs +++ b/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,7 +39,7 @@ 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). @@ -47,6 +47,10 @@ pub enum Source { /// A custom repodata source (provides records for requested platforms). Custom(Arc), + + /// A sparse repodata source (provides records for requested platforms from sparse + /// repodata). Each entry represents a different subdir. + SparseRepoData(Vec>), } impl From for Source { @@ -61,6 +65,18 @@ impl From> for Source { } } +impl From> for Source { + fn from(source: Arc) -> Self { + Source::SparseRepoData(vec![source]) + } +} + +impl From>> for Source { + fn from(sources: Vec>) -> Self { + Source::SparseRepoData(sources) + } +} + /// Adapts a [`RepoDataSource`] to the internal [`SubdirClient`] trait /// for a specific platform. /// diff --git a/crates/rattler_repodata_gateway/src/sparse/mod.rs b/crates/rattler_repodata_gateway/src/sparse/mod.rs index 5893f625c8..f66044a046 100644 --- a/crates/rattler_repodata_gateway/src/sparse/mod.rs +++ b/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, diff --git a/py-rattler/rattler/repo_data/gateway.py b/py-rattler/rattler/repo_data/gateway.py index 2986ef2759..a0975262d8 100644 --- a/py-rattler/rattler/repo_data/gateway.py +++ b/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." ) diff --git a/py-rattler/rattler/solver/solver.py b/py-rattler/rattler/solver/solver.py index d93abd21f9..fe06854da4 100644 --- a/py-rattler/rattler/solver/solver.py +++ b/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. diff --git a/py-rattler/src/repo_data/gateway.rs b/py-rattler/src/repo_data/gateway.rs index 970fdcfd77..794f4ae725 100644 --- a/py-rattler/src/repo_data/gateway.rs +++ b/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) -> 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 { @@ -154,6 +158,11 @@ pub fn py_object_to_source(obj: Bound<'_, PyAny>) -> PyResult { return Ok(Source::from(channel.inner)); } + // Then try to extract as SparseRepoData + if let Ok(sparse) = obj.extract::>() { + 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 { } 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 = Vec::new(); let mut custom_sources: Vec> = Vec::new(); + let mut sparse_sources: Vec> = + 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), } } diff --git a/py-rattler/src/repo_data/sparse.rs b/py-rattler/src/repo_data/sparse.rs index 96281ca38f..6432465b8a 100644 --- a/py-rattler/src/repo_data/sparse.rs +++ b/py-rattler/src/repo_data/sparse.rs @@ -22,10 +22,13 @@ 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>>, - 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>>>, + pub(crate) subdir: String, } impl PySparseRepoData { @@ -33,13 +36,24 @@ impl PySparseRepoData { pub(crate) fn from_args(channel: PyChannel, subdir: String, path: PathBuf) -> PyResult { 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> { + self.inner + .read() + .as_ref() + .cloned() + .ok_or_else(|| PyValueError::new_err("I/O operation on closed file.")) + } } impl From 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::, _>>()?; diff --git a/py-rattler/src/solver.rs b/py-rattler/src/solver.rs index 0e7713d50a..0933b477ea 100644 --- a/py-rattler/src/solver.rs +++ b/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::, _>>()?; diff --git a/py-rattler/tests/unit/test_solver.py b/py-rattler/tests/unit/test_solver.py index 0c68584545..e7e91b0f0b 100644 --- a/py-rattler/tests/unit/test_solver.py +++ b/py-rattler/tests/unit/test_solver.py @@ -164,6 +164,30 @@ async def test_solve_with_repodata() -> None: assert len(solved_data) == 2 +@pytest.mark.asyncio +async def test_solve_accepts_sparse_repodata_as_source() -> None: + """`solve` should accept `SparseRepoData` instances directly in `sources`, + without needing a `Gateway` to fetch them.""" + linux64_chan = Channel("conda-forge") + data_dir = os.path.join(os.path.dirname(__file__), "../../../test-data/") + linux64_path = os.path.join(data_dir, "channels/dummy/linux-64/repodata.json") + linux64_data = SparseRepoData( + channel=linux64_chan, + subdir="linux-64", + path=linux64_path, + ) + + solved_data = await solve( + [linux64_data], + ["foobar"], + platforms=["linux-64"], + ) + + assert isinstance(solved_data, list) + assert isinstance(solved_data[0], RepoDataRecord) + assert len(solved_data) == 2 + + @pytest.mark.asyncio async def test_conditional_root_requirement_satisfied(gateway: Gateway, dummy_channel: Channel) -> None: """Test that a conditional root requirement is included when the condition is satisfied."""