Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merge extras in lockfile #5181

Merged
merged 1 commit into from
Jul 18, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
88 changes: 44 additions & 44 deletions crates/uv-resolver/src/lock.rs
Original file line number Diff line number Diff line change
@@ -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};
Expand Down Expand Up @@ -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))
})
});
}
}
Expand Down Expand Up @@ -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)));
}
}
}
}
Expand Down Expand Up @@ -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
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This makes add_dependency quadratic, right?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I can't figure out a clean way around this.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is only called from Lock::from_resolution_graph and it's un-exported, so it should be fine to build up some intermediate state that replaces this loop with a O(1) or O(logn) check.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Totally agree. I briefly tried but had some trouble making it work.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For the large resolution transformers with all extras, this is the distribution of the number of iterations for the loop:

image

We run the code inside the loop 2762 times.

{
existing_dep.extra.extend(new_dep.extra);
return;
}
}
self.dependencies.push(new_dep);
}

/// Add the [`AnnotatedDist`] as an optional dependency of the [`Distribution`].
Expand Down Expand Up @@ -2030,7 +2028,7 @@ impl TryFrom<WheelWire> for Wheel {
#[derive(Clone, Debug, Eq, PartialEq, PartialOrd, Ord)]
struct Dependency {
distribution_id: DistributionId,
extra: Option<ExtraName>,
extra: BTreeSet<ExtraName>,
marker: Option<MarkerTree>,
}

Expand All @@ -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,
Expand All @@ -2056,13 +2054,11 @@ impl Dependency {
extras: &mut FxHashMap<PackageName, Vec<ExtraName>>,
) -> Result<Option<Requirement>, 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`.
Expand Down Expand Up @@ -2124,8 +2120,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::<Array>();
table.insert("extra", value(extra_array));
}
if let Some(ref marker) = self.marker {
table.insert("marker", value(marker.to_string()));
Expand All @@ -2137,20 +2138,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
)
Expand All @@ -2163,9 +2164,8 @@ impl std::fmt::Display for Dependency {
struct DependencyWire {
#[serde(flatten)]
distribution_id: DistributionIdForDependency,
#[serde(skip_serializing_if = "Option::is_none")]
extra: Option<ExtraName>,
#[serde(skip_serializing_if = "Option::is_none")]
#[serde(default)]
extra: BTreeSet<ExtraName>,
marker: Option<MarkerTree>,
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -149,7 +149,7 @@ Ok(
},
),
},
extra: None,
extra: {},
marker: None,
},
],
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -149,7 +149,7 @@ Ok(
},
),
},
extra: None,
extra: {},
marker: None,
},
],
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -149,7 +149,7 @@ Ok(
},
),
},
extra: None,
extra: {},
marker: None,
},
],
Expand Down
4 changes: 1 addition & 3 deletions crates/uv/tests/edit.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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]]
Expand Down
8 changes: 3 additions & 5 deletions crates/uv/tests/lock.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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]]
Expand Down Expand Up @@ -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]]
Expand Down Expand Up @@ -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]]
Expand Down
Loading