From 5bcdaedf8b0a5c5ccb2ede45824b34bc9dede8af Mon Sep 17 00:00:00 2001 From: konsti Date: Thu, 18 Jul 2024 20:07:49 +0200 Subject: [PATCH] Merge extras in lockfile (#5181) As user, you specify a list of extras. Internally, we decompose this into one virtual package per extra. We currently leak this abstraction by writing one entry per extra to the lockfile: ```toml [[distribution]] name = "foo" version = "4.39.0.dev0" source = { editable = "." } dependencies = [ { name = "pandas" }, { name = "pandas", extra = "excel" }, { name = "pandas", extra = "hdf5" }, { name = "pandas", extra = "html", marker = "os_name != 'posix'" }, { name = "pandas", extra = "output-formatting", marker = "os_name == 'posix'" }, { name = "pandas", extra = "plot", marker = "os_name == 'posix'" }, ] ``` Instead, we should merge the extras into a list of extras, creating a more concise lockfile: ```toml [[distribution]] name = "foo" version = "4.39.0.dev0" source = { editable = "." } dependencies = [ { name = "pandas", extra = ["excel", "hdf5"] }, { name = "pandas", extra = ["html"], marker = "os_name != 'posix'" }, { name = "pandas", extra = ["output-formatting", "plot"], marker = "os_name == 'posix'" }, ] ``` The base package is now implicitly included, as it is in PEP 508. Fixes #4888 --- crates/uv-resolver/src/lock.rs | 88 +++++++++---------- ...missing_dependency_source_unambiguous.snap | 2 +- ...dependency_source_version_unambiguous.snap | 2 +- ...issing_dependency_version_unambiguous.snap | 2 +- crates/uv/tests/edit.rs | 4 +- crates/uv/tests/lock.rs | 8 +- 6 files changed, 51 insertions(+), 55 deletions(-) diff --git a/crates/uv-resolver/src/lock.rs b/crates/uv-resolver/src/lock.rs index 8fed070975d3..7ce3c11b8bd8 100644 --- a/crates/uv-resolver/src/lock.rs +++ b/crates/uv-resolver/src/lock.rs @@ -1,13 +1,14 @@ #![allow(clippy::default_trait_access)] use std::borrow::Cow; -use std::collections::{BTreeMap, VecDeque}; +use std::collections::{BTreeMap, BTreeSet, VecDeque}; use std::fmt::{Debug, Display}; use std::path::{Path, PathBuf}; use std::str::FromStr; use std::sync::Arc; use either::Either; +use itertools::Itertools; use path_slash::PathExt; use petgraph::visit::EdgeRef; use rustc_hash::{FxHashMap, FxHashSet}; @@ -236,31 +237,16 @@ impl Lock { // Remove any non-existent extras (e.g., extras that were requested but don't exist). for dist in &mut distributions { - dist.dependencies.retain(|dep| { - dep.extra.as_ref().map_or(true, |extra| { + for dep in dist + .dependencies + .iter_mut() + .chain(dist.optional_dependencies.values_mut().flatten()) + .chain(dist.dev_dependencies.values_mut().flatten()) + { + dep.extra.retain(|extra| { extras_by_id .get(&dep.distribution_id) .is_some_and(|extras| extras.contains(extra)) - }) - }); - - for dependencies in dist.optional_dependencies.values_mut() { - dependencies.retain(|dep| { - dep.extra.as_ref().map_or(true, |extra| { - extras_by_id - .get(&dep.distribution_id) - .is_some_and(|extras| extras.contains(extra)) - }) - }); - } - - for dependencies in dist.dev_dependencies.values_mut() { - dependencies.retain(|dep| { - dep.extra.as_ref().map_or(true, |extra| { - extras_by_id - .get(&dep.distribution_id) - .is_some_and(|extras| extras.contains(extra)) - }) }); } } @@ -400,10 +386,14 @@ impl Lock { .as_ref() .map_or(true, |marker| marker.evaluate(marker_env, &[])) { - if seen.insert((&dep.distribution_id, dep.extra.as_ref())) { - let dep_dist = self.find_by_id(&dep.distribution_id); - let dep_extra = dep.extra.as_ref(); - queue.push_back((dep_dist, dep_extra)); + let dep_dist = self.find_by_id(&dep.distribution_id); + if seen.insert((&dep.distribution_id, None)) { + queue.push_back((dep_dist, None)); + } + for extra in &dep.extra { + if seen.insert((&dep.distribution_id, Some(extra))) { + queue.push_back((dep_dist, Some(extra))); + } } } } @@ -610,8 +600,16 @@ impl Distribution { /// Add the [`AnnotatedDist`] as a dependency of the [`Distribution`]. fn add_dependency(&mut self, annotated_dist: &AnnotatedDist, marker: Option<&MarkerTree>) { - self.dependencies - .push(Dependency::from_annotated_dist(annotated_dist, marker)); + let new_dep = Dependency::from_annotated_dist(annotated_dist, marker); + for existing_dep in &mut self.dependencies { + if existing_dep.distribution_id == new_dep.distribution_id + && existing_dep.marker == new_dep.marker + { + existing_dep.extra.extend(new_dep.extra); + return; + } + } + self.dependencies.push(new_dep); } /// Add the [`AnnotatedDist`] as an optional dependency of the [`Distribution`]. @@ -2030,7 +2028,7 @@ impl TryFrom for Wheel { #[derive(Clone, Debug, Eq, PartialEq, PartialOrd, Ord)] struct Dependency { distribution_id: DistributionId, - extra: Option, + extra: BTreeSet, marker: Option, } @@ -2040,7 +2038,7 @@ impl Dependency { marker: Option<&MarkerTree>, ) -> Dependency { let distribution_id = DistributionId::from_annotated_dist(annotated_dist); - let extra = annotated_dist.extra.clone(); + let extra = annotated_dist.extra.iter().cloned().collect(); let marker = marker.cloned(); Dependency { distribution_id, @@ -2056,13 +2054,11 @@ impl Dependency { extras: &mut FxHashMap>, ) -> Result, LockError> { // Keep track of extras, these will be denormalized later. - if let Some(ref extra) = self.extra { + if !self.extra.is_empty() { extras .entry(self.distribution_id.name.clone()) .or_default() - .push(extra.clone()); - - return Ok(None); + .extend(self.extra.iter().cloned()); } // Reconstruct the `RequirementSource` from the `Source`. @@ -2129,8 +2125,13 @@ impl Dependency { let mut table = Table::new(); self.distribution_id .to_toml(Some(dist_count_by_name), &mut table); - if let Some(ref extra) = self.extra { - table.insert("extra", value(extra.to_string())); + if !self.extra.is_empty() { + let extra_array = self + .extra + .iter() + .map(ToString::to_string) + .collect::(); + table.insert("extra", value(extra_array)); } if let Some(ref marker) = self.marker { table.insert("marker", value(marker.to_string())); @@ -2142,20 +2143,20 @@ impl Dependency { impl std::fmt::Display for Dependency { fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { - if let Some(ref extra) = self.extra { + if self.extra.is_empty() { write!( f, - "{}[{}]=={} @ {}", + "{}=={} @ {}", self.distribution_id.name, - extra, self.distribution_id.version, self.distribution_id.source ) } else { write!( f, - "{}=={} @ {}", + "{}[{}]=={} @ {}", self.distribution_id.name, + self.extra.iter().join(","), self.distribution_id.version, self.distribution_id.source ) @@ -2168,9 +2169,8 @@ impl std::fmt::Display for Dependency { struct DependencyWire { #[serde(flatten)] distribution_id: DistributionIdForDependency, - #[serde(skip_serializing_if = "Option::is_none")] - extra: Option, - #[serde(skip_serializing_if = "Option::is_none")] + #[serde(default)] + extra: BTreeSet, marker: Option, } diff --git a/crates/uv-resolver/src/snapshots/uv_resolver__lock__tests__missing_dependency_source_unambiguous.snap b/crates/uv-resolver/src/snapshots/uv_resolver__lock__tests__missing_dependency_source_unambiguous.snap index fe09635a4188..dc287f66e70c 100644 --- a/crates/uv-resolver/src/snapshots/uv_resolver__lock__tests__missing_dependency_source_unambiguous.snap +++ b/crates/uv-resolver/src/snapshots/uv_resolver__lock__tests__missing_dependency_source_unambiguous.snap @@ -149,7 +149,7 @@ Ok( }, ), }, - extra: None, + extra: {}, marker: None, }, ], diff --git a/crates/uv-resolver/src/snapshots/uv_resolver__lock__tests__missing_dependency_source_version_unambiguous.snap b/crates/uv-resolver/src/snapshots/uv_resolver__lock__tests__missing_dependency_source_version_unambiguous.snap index fe09635a4188..dc287f66e70c 100644 --- a/crates/uv-resolver/src/snapshots/uv_resolver__lock__tests__missing_dependency_source_version_unambiguous.snap +++ b/crates/uv-resolver/src/snapshots/uv_resolver__lock__tests__missing_dependency_source_version_unambiguous.snap @@ -149,7 +149,7 @@ Ok( }, ), }, - extra: None, + extra: {}, marker: None, }, ], diff --git a/crates/uv-resolver/src/snapshots/uv_resolver__lock__tests__missing_dependency_version_unambiguous.snap b/crates/uv-resolver/src/snapshots/uv_resolver__lock__tests__missing_dependency_version_unambiguous.snap index fe09635a4188..dc287f66e70c 100644 --- a/crates/uv-resolver/src/snapshots/uv_resolver__lock__tests__missing_dependency_version_unambiguous.snap +++ b/crates/uv-resolver/src/snapshots/uv_resolver__lock__tests__missing_dependency_version_unambiguous.snap @@ -149,7 +149,7 @@ Ok( }, ), }, - extra: None, + extra: {}, marker: None, }, ], diff --git a/crates/uv/tests/edit.rs b/crates/uv/tests/edit.rs index 26c2679beab1..9baeafdaf36f 100644 --- a/crates/uv/tests/edit.rs +++ b/crates/uv/tests/edit.rs @@ -1405,9 +1405,7 @@ fn update() -> Result<()> { version = "0.1.0" source = { editable = "." } dependencies = [ - { name = "requests" }, - { name = "requests", extra = "socks" }, - { name = "requests", extra = "use-chardet-on-py3" }, + { name = "requests", extra = ["socks", "use-chardet-on-py3"] }, ] [[distribution]] diff --git a/crates/uv/tests/lock.rs b/crates/uv/tests/lock.rs index b380cf5456a6..d4ca1846ca56 100644 --- a/crates/uv/tests/lock.rs +++ b/crates/uv/tests/lock.rs @@ -776,8 +776,7 @@ fn lock_dependency_extra() -> Result<()> { version = "0.1.0" source = { editable = "." } dependencies = [ - { name = "flask" }, - { name = "flask", extra = "dotenv" }, + { name = "flask", extra = ["dotenv"] }, ] [[distribution]] @@ -989,7 +988,7 @@ fn lock_conditional_dependency_extra() -> Result<()> { source = { editable = "." } dependencies = [ { name = "requests" }, - { name = "requests", extra = "socks", marker = "python_version < '3.10'" }, + { name = "requests", extra = ["socks"], marker = "python_version < '3.10'" }, ] [[distribution]] @@ -3872,8 +3871,7 @@ fn lock_new_extras() -> Result<()> { version = "0.1.0" source = { editable = "." } dependencies = [ - { name = "requests" }, - { name = "requests", extra = "socks" }, + { name = "requests", extra = ["socks"] }, ] [[distribution]]