diff --git a/crates/uv-resolver/src/lock/tree.rs b/crates/uv-resolver/src/lock/tree.rs index cd388545cb8..9dbe953d9c7 100644 --- a/crates/uv-resolver/src/lock/tree.rs +++ b/crates/uv-resolver/src/lock/tree.rs @@ -1,3 +1,4 @@ +use std::cmp::Ordering; use std::collections::{BTreeSet, VecDeque}; use std::fmt::Write; @@ -146,7 +147,11 @@ impl<'env> TreeDisplay<'env> { .or_insert_with(|| graph.add_node(Node::Package(&dep.package_id))); // Add an edge from the workspace package. - graph.add_edge(index, dep_index, Edge::Dev(group, Some(&dep.extra))); + graph.add_edge( + index, + dep_index, + Edge::Dev(group, Some(RequestedExtras::Dependency(&dep.extra))), + ); // Push its dependencies on the queue. if seen.insert((&dep.package_id, None)) { @@ -207,17 +212,31 @@ impl<'env> TreeDisplay<'env> { .or_insert_with(|| graph.add_node(Node::Package(&package.id))); // Add an edge from the root. - graph.add_edge(root, *index, Edge::Prod(None)); + graph.add_edge( + root, + *index, + Edge::Prod(Some(RequestedExtras::Requirement( + requirement.extras.as_ref(), + ))), + ); // Push its dependencies on the queue. if seen.insert((&package.id, None)) { queue.push_back((&package.id, None)); } + for extra in &*requirement.extras { + if seen.insert((&package.id, Some(extra))) { + queue.push_back((&package.id, Some(extra))); + } + } } } // Identify any dependency groups attached to the workspace itself. for (group, requirements) in lock.dependency_groups() { + if !groups.contains(group) { + continue; + } for requirement in requirements { for package in by_name.get(&requirement.name).into_iter().flatten() { // Determine whether this entry is "relevant" for the requirement, by intersecting @@ -244,12 +263,24 @@ impl<'env> TreeDisplay<'env> { .or_insert_with(|| graph.add_node(Node::Package(&package.id))); // Add an edge from the root. - graph.add_edge(root, *index, Edge::Dev(group, None)); + graph.add_edge( + root, + *index, + Edge::Dev( + group, + Some(RequestedExtras::Requirement(requirement.extras.as_ref())), + ), + ); // Push its dependencies on the queue. if seen.insert((&package.id, None)) { queue.push_back((&package.id, None)); } + for extra in &*requirement.extras { + if seen.insert((&package.id, Some(extra))) { + queue.push_back((&package.id, Some(extra))); + } + } } } } @@ -293,9 +324,9 @@ impl<'env> TreeDisplay<'env> { index, dep_index, if let Some(extra) = extra { - Edge::Optional(extra, Some(&dep.extra)) + Edge::Optional(extra, Some(RequestedExtras::Dependency(&dep.extra))) } else { - Edge::Prod(Some(&dep.extra)) + Edge::Prod(Some(RequestedExtras::Dependency(&dep.extra))) }, ); @@ -665,17 +696,15 @@ enum Node<'env> { #[derive(Debug, Clone, PartialEq, Eq, Ord, PartialOrd)] enum Edge<'env> { - Prod(Option<&'env BTreeSet>), - Optional(&'env ExtraName, Option<&'env BTreeSet>), - Dev(&'env GroupName, Option<&'env BTreeSet>), + Prod(Option>), + Optional(&'env ExtraName, Option>), + Dev(&'env GroupName, Option>), } impl<'env> Edge<'env> { - fn extras(&self) -> Option<&'env BTreeSet> { + fn extras(&self) -> Option> { match self { - Self::Prod(extras) => *extras, - Self::Optional(_, extras) => *extras, - Self::Dev(_, extras) => *extras, + Self::Prod(extras) | Self::Optional(_, extras) | Self::Dev(_, extras) => *extras, } } @@ -688,6 +717,48 @@ impl<'env> Edge<'env> { } } +#[derive(Debug, Copy, Clone)] +enum RequestedExtras<'env> { + Dependency(&'env BTreeSet), + Requirement(&'env [ExtraName]), +} + +impl PartialEq for RequestedExtras<'_> { + fn eq(&self, other: &Self) -> bool { + self.iter().eq(other.iter()) + } +} + +impl Eq for RequestedExtras<'_> {} + +impl PartialOrd for RequestedExtras<'_> { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl Ord for RequestedExtras<'_> { + fn cmp(&self, other: &Self) -> Ordering { + self.iter().cmp(other.iter()) + } +} + +impl<'env> RequestedExtras<'env> { + fn is_empty(self) -> bool { + match self { + Self::Dependency(extras) => extras.is_empty(), + Self::Requirement(extras) => extras.is_empty(), + } + } + + fn iter(self) -> impl Iterator { + match self { + Self::Dependency(extras) => Either::Left(extras.iter()), + Self::Requirement(extras) => Either::Right(extras.iter()), + } + } +} + #[derive(Debug, Clone, PartialEq, Eq, Ord, PartialOrd)] enum EdgeKind<'env> { Prod, diff --git a/crates/uv/tests/project/tree.rs b/crates/uv/tests/project/tree.rs index 0dcaf862876..7555e2119de 100644 --- a/crates/uv/tests/project/tree.rs +++ b/crates/uv/tests/project/tree.rs @@ -1655,7 +1655,7 @@ fn non_project() -> Result<()> { "#, )?; - uv_snapshot!(context.filters(), context.tree().arg("--universal"), @" + uv_snapshot!(context.filters(), context.tree().arg("--universal").arg("--group").arg("async"), @" success: true exit_code: 0 ----- stdout ----- @@ -1676,6 +1676,98 @@ fn non_project() -> Result<()> { Ok(()) } +#[test] +fn non_project_group_selection_with_extras() -> Result<()> { + let context = uv_test::test_context!("3.12"); + + let leaf = context.temp_dir.child("leaf"); + leaf.create_dir_all()?; + leaf.child("pyproject.toml").write_str( + r#" + [project] + name = "leaf" + version = "0.1.0" + requires-python = ">=3.12" + "#, + )?; + let leaf_url = Url::from_file_path(leaf.path()) + .map_err(|()| anyhow::anyhow!("failed to convert leaf path to URL"))?; + + let child = context.temp_dir.child("child"); + child.create_dir_all()?; + child.child("pyproject.toml").write_str(&formatdoc! { + r#" + [project] + name = "child" + version = "0.1.0" + requires-python = ">=3.12" + + [project.optional-dependencies] + feature = ["leaf @ {leaf_url}"] + "#, + })?; + let child_url = Url::from_file_path(child.path()) + .map_err(|()| anyhow::anyhow!("failed to convert child path to URL"))?; + + let test_dependency = context.temp_dir.child("test-dependency"); + test_dependency.create_dir_all()?; + test_dependency.child("pyproject.toml").write_str( + r#" + [project] + name = "test-dependency" + version = "0.1.0" + requires-python = ">=3.12" + "#, + )?; + let test_dependency_url = Url::from_file_path(test_dependency.path()) + .map_err(|()| anyhow::anyhow!("failed to convert test dependency path to URL"))?; + + let pyproject_toml = context.temp_dir.child("pyproject.toml"); + pyproject_toml.write_str(&formatdoc! { + r#" + [tool.uv.workspace] + members = [] + + [dependency-groups] + dev = ["child[feature] @ {child_url}"] + test = ["test-dependency @ {test_dependency_url}"] + "#, + })?; + + uv_snapshot!(context.filters(), context.tree().arg("--only-group").arg("dev"), @" + success: true + exit_code: 0 + ----- stdout ----- + child[feature] v0.1.0 (group: dev) + └── leaf v0.1.0 (extra: feature) + + ----- stderr ----- + warning: No `requires-python` value found in the workspace. Defaulting to `>=3.12`. + Resolved 3 packages in [TIME] + "); + + let script = context.temp_dir.child("script.py"); + script.write_str(&formatdoc! {r#" + # /// script + # requires-python = ">=3.12" + # dependencies = ["child[feature] @ {child_url}"] + # /// + "#})?; + + uv_snapshot!(context.filters(), context.tree().arg("--script").arg(script.path()), @" + success: true + exit_code: 0 + ----- stdout ----- + child[feature] v0.1.0 + └── leaf v0.1.0 (extra: feature) + + ----- stderr ----- + Resolved 2 packages in [TIME] + "); + + Ok(()) +} + #[test] fn non_project_member() -> Result<()> { let context = uv_test::test_context!("3.12"); @@ -1706,7 +1798,7 @@ fn non_project_member() -> Result<()> { "#, )?; - uv_snapshot!(context.filters(), context.tree().arg("--universal"), @" + uv_snapshot!(context.filters(), context.tree().arg("--universal").arg("--group").arg("async"), @" success: true exit_code: 0 ----- stdout ----- @@ -1724,7 +1816,7 @@ fn non_project_member() -> Result<()> { " ); - uv_snapshot!(context.filters(), context.tree().arg("--universal").arg("--invert"), @" + uv_snapshot!(context.filters(), context.tree().arg("--universal").arg("--invert").arg("--group").arg("async"), @" success: true exit_code: 0 ----- stdout -----