Add extra dependency annotations to lockfile and sync commands (#3913)

## Summary

This PR adds extras to the lockfile, and enables users to selectively
sync extras in `uv sync` and `uv run`. The end result here was fairly
simple, though it required a few refactors to get here. The basic idea
is that `DistributionId` now includes `extra: Option<ExtraName>`, so we
effectively treat extras as separate packages. Generating the lockfile,
and generating the resolution from the lockfile, fall out of this
naturally with no special-casing or additional changes.

The main downside here is that it bloats the lockfile significantly.
Specifically:

- We include _all_ distribution URLs and hashes for _every_ extra
variant.
- We include all dependencies for the extra variant, even though that
are dependencies of the base package.

We could normalize this representation by changing each distribution
have an `optional-dependencies` hash map that keys on extras, but we
actually don't have the information we need to create that right now
(specifically, we can't differentiate between dependencies that
_require_ the extra and dependencies on the base package).

Closes #3700.
This commit is contained in:
Charlie Marsh 2024-05-29 15:25:58 -04:00 committed by GitHub
parent 1bd5d8bc34
commit 4859a27948
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 142 additions and 98 deletions

View file

@ -20,7 +20,7 @@ use pep508_rs::{MarkerEnvironment, VerbatimUrl};
use platform_tags::{TagCompatibility, TagPriority, Tags};
use pypi_types::{HashDigest, ParsedArchiveUrl, ParsedGitUrl};
use uv_git::{GitReference, GitSha};
use uv_normalize::PackageName;
use uv_normalize::{ExtraName, PackageName};
use crate::resolution::AnnotatedDist;
@ -57,18 +57,18 @@ impl Lock {
marker_env: &MarkerEnvironment,
tags: &Tags,
root_name: &PackageName,
extras: &[ExtraName],
) -> Resolution {
let root = self
.find_by_name(root_name)
// TODO: In the future, we should derive the root distribution
// from the pyproject.toml, but I don't think the infrastructure
// for that is in place yet. For now, we ask the caller to specify
// the root package name explicitly, and we assume here that it is
// correct.
.expect("found too many distributions matching root")
.expect("could not find root");
let mut queue: VecDeque<&Distribution> = VecDeque::new();
queue.push_back(root);
// Add the root distribution to the queue.
for extra in std::iter::once(None).chain(extras.iter().map(Some)) {
let root = self
.find_by_name(root_name, extra)
.expect("found too many distributions matching root")
.expect("could not find root");
queue.push_back(root);
}
let mut map = BTreeMap::default();
while let Some(dist) = queue.pop_front() {
@ -87,12 +87,20 @@ impl Lock {
/// Returns the distribution with the given name. If there are multiple
/// matching distributions, then an error is returned. If there are no
/// matching distributions, then `Ok(None)` is returned.
fn find_by_name(&self, name: &PackageName) -> Result<Option<&Distribution>, String> {
fn find_by_name(
&self,
name: &PackageName,
extra: Option<&ExtraName>,
) -> Result<Option<&Distribution>, String> {
let mut found_dist = None;
for dist in &self.distributions {
if &dist.id.name == name {
if &dist.id.name == name && dist.id.extra.as_ref() == extra {
if found_dist.is_some() {
return Err(format!("found multiple distributions matching `{name}`"));
return Err(if let Some(extra) = extra {
format!("found multiple distributions matching `{name}[{extra}]`")
} else {
format!("found multiple distributions matching `{name}`")
});
}
found_dist = Some(dist);
}
@ -407,6 +415,8 @@ impl Distribution {
pub(crate) struct DistributionId {
pub(crate) name: PackageName,
pub(crate) version: Version,
#[serde(skip_serializing_if = "Option::is_none")]
pub(crate) extra: Option<ExtraName>,
pub(crate) source: Source,
}
@ -414,10 +424,12 @@ impl DistributionId {
fn from_annotated_dist(annotated_dist: &AnnotatedDist) -> DistributionId {
let name = annotated_dist.metadata.name.clone();
let version = annotated_dist.metadata.version.clone();
let extra = annotated_dist.extra.clone();
let source = Source::from_resolved_dist(&annotated_dist.dist);
DistributionId {
name,
version,
extra,
source,
}
}

View file

@ -47,7 +47,11 @@ impl ResolutionGraph {
preferences: &Preferences,
) -> anyhow::Result<Self, ResolveError> {
// Add every package to the graph.
let mut petgraph = petgraph::graph::Graph::with_capacity(selection.len(), selection.len());
let mut petgraph =
petgraph::graph::Graph::<AnnotatedDist, (), petgraph::Directed>::with_capacity(
selection.len(),
selection.len(),
);
let mut inverse =
FxHashMap::with_capacity_and_hasher(selection.len(), BuildHasherDefault::default());
let mut diagnostics = Vec::new();
@ -248,23 +252,30 @@ impl ResolutionGraph {
continue;
};
let PubGrubPackageInner::Package {
name: dependency_name,
extra: dependency_extra,
..
} = &**dependency_package
else {
continue;
};
match &**dependency_package {
PubGrubPackageInner::Package {
name: dependency_name,
extra: dependency_extra,
..
} => {
let self_index = &inverse[&(self_name, self_extra)];
let dependency_index = &inverse[&(dependency_name, dependency_extra)];
petgraph.update_edge(*self_index, *dependency_index, ());
}
// For extras, we include a dependency between the extra and the base package.
if self_name == dependency_name {
continue;
PubGrubPackageInner::Extra {
name: dependency_name,
extra: dependency_extra,
..
} => {
let self_index = &inverse[&(self_name, self_extra)];
let dependency_extra = Some(dependency_extra.clone());
let dependency_index = &inverse[&(dependency_name, &dependency_extra)];
petgraph.update_edge(*self_index, *dependency_index, ());
}
_ => {}
}
let self_index = &inverse[&(self_name, self_extra)];
let dependency_index = &inverse[&(dependency_name, dependency_extra)];
petgraph.update_edge(*self_index, *dependency_index, ());
}
}
}
@ -445,10 +456,6 @@ impl ResolutionGraph {
let mut locked_dists = vec![];
for node_index in self.petgraph.node_indices() {
let dist = &self.petgraph[node_index];
if dist.extra.is_some() {
continue;
}
let mut locked_dist = lock::Distribution::from_annotated_dist(dist)?;
for edge in self.petgraph.neighbors(node_index) {
let dependency_dist = &self.petgraph[edge];

View file

@ -853,7 +853,6 @@ impl<InstalledPackages: InstalledPackagesProvider> ResolverState<InstalledPackag
if self.dependency_mode.is_direct() {
// If an extra is provided, wait for the metadata to be available, since it's
// still required for generating the lock file.
let dist = match url {
Some(url) => PubGrubDistribution::from_url(name, url),
None => PubGrubDistribution::from_registry(name, version),

View file

@ -12,6 +12,7 @@ Ok(
"anyio",
),
version: "4.3.0",
extra: None,
source: Source {
kind: Path,
url: Url {
@ -78,6 +79,7 @@ Ok(
"anyio",
),
version: "4.3.0",
extra: None,
source: Source {
kind: Path,
url: Url {

View file

@ -1764,6 +1764,17 @@ pub(crate) struct VenvArgs {
#[derive(Args)]
#[allow(clippy::struct_excessive_bools)]
pub(crate) struct RunArgs {
/// Include optional dependencies in the given extra group name; may be provided more than once.
#[arg(long, conflicts_with = "all_extras", value_parser = extra_name_with_clap_error)]
pub(crate) extra: Option<Vec<ExtraName>>,
/// Include all optional dependencies.
#[arg(long, conflicts_with = "extra")]
pub(crate) all_extras: bool,
#[arg(long, overrides_with("all_extras"), hide = true)]
pub(crate) no_all_extras: bool,
/// The command to run.
pub(crate) target: Option<String>,
@ -1800,6 +1811,17 @@ pub(crate) struct RunArgs {
#[derive(Args)]
#[allow(clippy::struct_excessive_bools)]
pub(crate) struct SyncArgs {
/// Include optional dependencies in the given extra group name; may be provided more than once.
#[arg(long, conflicts_with = "all_extras", value_parser = extra_name_with_clap_error)]
pub(crate) extra: Option<Vec<ExtraName>>,
/// Include all optional dependencies.
#[arg(long, conflicts_with = "extra")]
pub(crate) all_extras: bool,
#[arg(long, overrides_with("all_extras"), hide = true)]
pub(crate) no_all_extras: bool,
/// The Python interpreter to use to build the run environment.
///
/// By default, `uv` uses the virtual environment in the current working directory or any parent

View file

@ -323,7 +323,7 @@ pub(crate) async fn pip_install(
let root = PackageName::new(root.to_string())?;
let encoded = fs::tokio::read_to_string("uv.lock").await?;
let lock: Lock = toml::from_str(&encoded)?;
lock.to_resolution(&markers, &tags, &root)
lock.to_resolution(&markers, &tags, &root, &[])
} else {
let options = OptionsBuilder::new()
.resolution_mode(resolution_mode)

View file

@ -9,7 +9,7 @@ use tracing::debug;
use uv_cache::Cache;
use uv_client::Connectivity;
use uv_configuration::PreviewMode;
use uv_configuration::{ExtrasSpecification, PreviewMode};
use uv_interpreter::{PythonEnvironment, SystemPython};
use uv_requirements::{ProjectWorkspace, RequirementsSource};
use uv_resolver::ExcludeNewer;
@ -21,6 +21,7 @@ use crate::printer::Printer;
/// Run a command.
#[allow(clippy::too_many_arguments)]
pub(crate) async fn run(
extras: ExtrasSpecification,
target: Option<String>,
mut args: Vec<OsString>,
requirements: Vec<RequirementsSource>,
@ -47,7 +48,7 @@ pub(crate) async fn run(
// Lock and sync the environment.
let lock = project::lock::do_lock(&project, &venv, exclude_newer, cache, printer).await?;
project::sync::do_sync(&project, &venv, &lock, cache, printer).await?;
project::sync::do_sync(&project, &venv, &lock, extras, cache, printer).await?;
Some(venv)
};

View file

@ -5,7 +5,8 @@ use install_wheel_rs::linker::LinkMode;
use uv_cache::Cache;
use uv_client::RegistryClientBuilder;
use uv_configuration::{
Concurrency, ConfigSettings, NoBinary, NoBuild, PreviewMode, Reinstall, SetupPyStrategy,
Concurrency, ConfigSettings, ExtrasSpecification, NoBinary, NoBuild, PreviewMode, Reinstall,
SetupPyStrategy,
};
use uv_dispatch::BuildDispatch;
use uv_installer::SitePackages;
@ -23,6 +24,7 @@ use crate::printer::Printer;
/// Sync the project environment.
#[allow(clippy::too_many_arguments)]
pub(crate) async fn sync(
extras: ExtrasSpecification,
preview: PreviewMode,
cache: &Cache,
printer: Printer,
@ -45,7 +47,7 @@ pub(crate) async fn sync(
};
// Perform the sync operation.
do_sync(&project, &venv, &lock, cache, printer).await?;
do_sync(&project, &venv, &lock, extras, cache, printer).await?;
Ok(ExitStatus::Success)
}
@ -55,6 +57,7 @@ pub(super) async fn do_sync(
project: &ProjectWorkspace,
venv: &PythonEnvironment,
lock: &Lock,
extras: ExtrasSpecification,
cache: &Cache,
printer: Printer,
) -> Result<(), ProjectError> {
@ -62,7 +65,12 @@ pub(super) async fn do_sync(
let tags = venv.interpreter().tags()?;
// Read the lockfile.
let resolution = lock.to_resolution(markers, tags, project.project_name());
let extras = match extras {
ExtrasSpecification::None => vec![],
ExtrasSpecification::All => project.project_extras().to_vec(),
ExtrasSpecification::Some(extras) => extras,
};
let resolution = lock.to_resolution(markers, tags, project.project_name(), &extras);
// Initialize the registry client.
// TODO(zanieb): Support client options e.g. offline, tls, etc.

View file

@ -573,6 +573,7 @@ async fn run() -> Result<ExitStatus> {
.collect::<Vec<_>>();
commands::run(
args.extras,
args.target,
args.args,
requirements,
@ -588,12 +589,12 @@ async fn run() -> Result<ExitStatus> {
}
Commands::Sync(args) => {
// Resolve the settings from the command-line arguments and workspace configuration.
let _args = settings::SyncSettings::resolve(args, workspace);
let args = settings::SyncSettings::resolve(args, workspace);
// Initialize the cache.
let cache = cache.init()?;
commands::sync(globals.preview, &cache, printer).await
commands::sync(args.extras, globals.preview, &cache, printer).await
}
Commands::Lock(args) => {
// Resolve the settings from the command-line arguments and workspace configuration.

View file

@ -97,6 +97,7 @@ impl CacheSettings {
#[allow(clippy::struct_excessive_bools)]
#[derive(Debug, Clone)]
pub(crate) struct RunSettings {
pub(crate) extras: ExtrasSpecification,
pub(crate) target: Option<String>,
pub(crate) args: Vec<OsString>,
pub(crate) with: Vec<String>,
@ -109,6 +110,9 @@ impl RunSettings {
#[allow(clippy::needless_pass_by_value)]
pub(crate) fn resolve(args: RunArgs, _workspace: Option<Workspace>) -> Self {
let RunArgs {
extra,
all_extras,
no_all_extras,
target,
args,
with,
@ -117,6 +121,10 @@ impl RunSettings {
} = args;
Self {
extras: ExtrasSpecification::from_args(
flag(all_extras, no_all_extras).unwrap_or_default(),
extra.unwrap_or_default(),
),
target,
args,
with,
@ -130,6 +138,7 @@ impl RunSettings {
#[allow(clippy::struct_excessive_bools, dead_code)]
#[derive(Debug, Clone)]
pub(crate) struct SyncSettings {
pub(crate) extras: ExtrasSpecification,
pub(crate) python: Option<String>,
}
@ -137,9 +146,20 @@ impl SyncSettings {
/// Resolve the [`SyncSettings`] from the CLI and workspace configuration.
#[allow(clippy::needless_pass_by_value)]
pub(crate) fn resolve(args: SyncArgs, _workspace: Option<Workspace>) -> Self {
let SyncArgs { python } = args;
let SyncArgs {
extra,
all_extras,
no_all_extras,
python,
} = args;
Self { python }
Self {
extras: ExtrasSpecification::from_args(
flag(all_extras, no_all_extras).unwrap_or_default(),
extra.unwrap_or_default(),
),
python,
}
}
}

View file

@ -576,7 +576,7 @@ fn lock_extra() -> Result<()> {
dependencies = ["anyio==3.7.0"]
[project.optional-dependencies]
test = ["pytest"]
test = ["iniconfig"]
"#,
)?;
@ -587,7 +587,7 @@ fn lock_extra() -> Result<()> {
----- stderr -----
warning: `uv lock` is experimental and may change without warning.
Resolved 8 packages in [TIME]
Resolved 5 packages in [TIME]
"###);
let lock = fs_err::read_to_string(context.temp_dir.join("uv.lock"))?;
@ -654,36 +654,6 @@ fn lock_extra() -> Result<()> {
hash = "sha256:b6a85871a79d2e3b22d2d1b94ac2824226a63c6b741c88f7ae975f18b6778374"
size = 5892
[[distribution]]
name = "packaging"
version = "24.0"
source = "registry+https://pypi.org/simple"
[distribution.sdist]
url = "https://files.pythonhosted.org/packages/ee/b5/b43a27ac7472e1818c4bafd44430e69605baefe1f34440593e0332ec8b4d/packaging-24.0.tar.gz"
hash = "sha256:eb82c5e3e56209074766e6885bb04b8c38a0c015d0a30036ebe7ece34c9989e9"
size = 147882
[[distribution.wheel]]
url = "https://files.pythonhosted.org/packages/49/df/1fceb2f8900f8639e278b056416d49134fb8d84c5942ffaa01ad34782422/packaging-24.0-py3-none-any.whl"
hash = "sha256:2ddfb553fdf02fb784c234c7ba6ccc288296ceabec964ad2eae3777778130bc5"
size = 53488
[[distribution]]
name = "pluggy"
version = "1.4.0"
source = "registry+https://pypi.org/simple"
[distribution.sdist]
url = "https://files.pythonhosted.org/packages/54/c6/43f9d44d92aed815e781ca25ba8c174257e27253a94630d21be8725a2b59/pluggy-1.4.0.tar.gz"
hash = "sha256:8c85c2876142a764e5b7548e7d9a0e0ddb46f5185161049a79b7e974454223be"
size = 65812
[[distribution.wheel]]
url = "https://files.pythonhosted.org/packages/a5/5b/0cc789b59e8cc1bf288b38111d002d8c5917123194d45b29dcdac64723cc/pluggy-1.4.0-py3-none-any.whl"
hash = "sha256:7db9f7b503d67d1c5b95f59773ebb58a8c1c288129a88665838012cfb07b8981"
size = 20120
[[distribution]]
name = "project"
version = "0.1.0"
@ -698,35 +668,24 @@ fn lock_extra() -> Result<()> {
source = "registry+https://pypi.org/simple"
[[distribution]]
name = "pytest"
version = "8.1.1"
source = "registry+https://pypi.org/simple"
name = "project"
version = "0.1.0"
extra = "test"
source = "editable+file://[TEMP_DIR]/"
[distribution.sdist]
url = "https://files.pythonhosted.org/packages/30/b7/7d44bbc04c531dcc753056920e0988032e5871ac674b5a84cb979de6e7af/pytest-8.1.1.tar.gz"
hash = "sha256:ac978141a75948948817d360297b7aae0fcb9d6ff6bc9ec6d514b85d5a65c044"
size = 1409703
url = "file://[TEMP_DIR]/"
[[distribution.wheel]]
url = "https://files.pythonhosted.org/packages/4d/7e/c79cecfdb6aa85c6c2e3cf63afc56d0f165f24f5c66c03c695c4d9b84756/pytest-8.1.1-py3-none-any.whl"
hash = "sha256:2a8386cfc11fa9d2c50ee7b2a57e7d898ef90470a7a34c4b949ff59662bb78b7"
size = 337359
[[distribution.dependencies]]
name = "anyio"
version = "3.7.0"
source = "registry+https://pypi.org/simple"
[[distribution.dependencies]]
name = "iniconfig"
version = "2.0.0"
source = "registry+https://pypi.org/simple"
[[distribution.dependencies]]
name = "packaging"
version = "24.0"
source = "registry+https://pypi.org/simple"
[[distribution.dependencies]]
name = "pluggy"
version = "1.4.0"
source = "registry+https://pypi.org/simple"
[[distribution]]
name = "sniffio"
version = "1.3.1"
@ -745,7 +704,7 @@ fn lock_extra() -> Result<()> {
);
});
// Install from the lockfile.
// Install the base dependencies from the lockfile.
uv_snapshot!(context.filters(), context.sync(), @r###"
success: true
exit_code: 0
@ -761,5 +720,18 @@ fn lock_extra() -> Result<()> {
+ sniffio==1.3.1
"###);
// Install the extras from the lockfile.
uv_snapshot!(context.filters(), context.sync().arg("--extra").arg("test"), @r###"
success: true
exit_code: 0
----- stdout -----
----- stderr -----
warning: `uv sync` is experimental and may change without warning.
Downloaded 1 package in [TIME]
Installed 1 package in [TIME]
+ iniconfig==2.0.0
"###);
Ok(())
}