Allow multiple pinned indexes in tool.uv.sources (#7769)

## Summary

This PR lifts the restriction that a package must come from a single
index. For example, you can now do:

```toml
[project]
name = "project"
version = "0.1.0"
readme = "README.md"
requires-python = ">=3.12"
dependencies = ["jinja2"]

[tool.uv.sources]
jinja2 = [
    { index = "torch-cu118", marker = "sys_platform == 'darwin'"},
    { index = "torch-cu124", marker = "sys_platform != 'darwin'"},
]

[[tool.uv.index]]
name = "torch-cu118"
url = "https://download.pytorch.org/whl/cu118"

[[tool.uv.index]]
name = "torch-cu124"
url = "https://download.pytorch.org/whl/cu124"
```

The construction is very similar to the way we handle URLs today: you
can have multiple URLs for a given package, but they must appear in
disjoint forks. So most of the code is just adding that abstraction to
the resolver, following our handling of URLs.

Closes #7761.
This commit is contained in:
Charlie Marsh 2024-10-15 15:58:15 -07:00 committed by GitHub
parent ad24cee7c6
commit 9a76e47888
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
14 changed files with 715 additions and 109 deletions

View file

@ -6,13 +6,12 @@ use rustc_hash::FxHashMap;
use tokio::sync::mpsc::Sender;
use tracing::{debug, trace};
use uv_distribution_types::{CompatibleDist, DistributionMetadata, IndexCapabilities};
use uv_pep440::Version;
use crate::candidate_selector::CandidateSelector;
use crate::pubgrub::{PubGrubPackage, PubGrubPackageInner};
use crate::resolver::Request;
use crate::{InMemoryIndex, PythonRequirement, ResolveError, ResolverMarkers, VersionsResponse};
use uv_distribution_types::{CompatibleDist, DistributionMetadata, IndexCapabilities, IndexUrl};
use uv_pep440::Version;
enum BatchPrefetchStrategy {
/// Go through the next versions assuming the existing selection and its constraints
@ -47,11 +46,12 @@ impl BatchPrefetcher {
pub(crate) fn prefetch_batches(
&mut self,
next: &PubGrubPackage,
index: Option<&IndexUrl>,
version: &Version,
current_range: &Range<Version>,
python_requirement: &PythonRequirement,
request_sink: &Sender<Request>,
index: &InMemoryIndex,
in_memory: &InMemoryIndex,
capabilities: &IndexCapabilities,
selector: &CandidateSelector,
markers: &ResolverMarkers,
@ -73,10 +73,17 @@ impl BatchPrefetcher {
let total_prefetch = min(num_tried, 50);
// This is immediate, we already fetched the version map.
let versions_response = index
.packages()
.wait_blocking(name)
.ok_or_else(|| ResolveError::UnregisteredTask(name.to_string()))?;
let versions_response = if let Some(index) = index {
in_memory
.explicit()
.wait_blocking(&(name.clone(), index.clone()))
.ok_or_else(|| ResolveError::UnregisteredTask(name.to_string()))?
} else {
in_memory
.implicit()
.wait_blocking(name)
.ok_or_else(|| ResolveError::UnregisteredTask(name.to_string()))?
};
let VersionsResponse::Found(ref version_map) = *versions_response else {
return Ok(());
@ -191,7 +198,7 @@ impl BatchPrefetcher {
);
prefetch_count += 1;
if index.distributions().register(candidate.version_id()) {
if in_memory.distributions().register(candidate.version_id()) {
let request = Request::from(dist);
request_sink.blocking_send(request)?;
}