Simplify selected extra markers in uv export

crates/uv-resolver/src/lock/export/mod.rs

@@ -55,6 +55,8 @@ impl<'lock> ExportableRequirements<'lock> {
         let size_guess = target.lock().packages.len();
         let mut graph = Graph::<Node<'lock>, Edge<'lock>>::with_capacity(size_guess, size_guess);
         let mut inverse = FxHashMap::with_capacity_and_hasher(size_guess, FxBuildHasher);
+        let mut selected_extras: FxHashMap<_, Vec<ExtraName>> =
+            FxHashMap::with_capacity_and_hasher(size_guess, FxBuildHasher);
 
         let mut queue: VecDeque<(&Package, Option<&ExtraName>)> = VecDeque::new();
         let mut seen = FxHashSet::default();
@@ -91,7 +93,9 @@ impl<'lock> ExportableRequirements<'lock> {
 
                 // Push its dependencies on the queue.
                 queue.push_back((dist, None));
+                let mut root_extras = Vec::new();
                 for extra in extras.extra_names(dist.optional_dependencies.keys()) {
+                    root_extras.push(extra.clone());
                     queue.push_back((dist, Some(extra)));
 
                     // Track the activated extra in the list of known conflicts.
@@ -102,6 +106,9 @@ impl<'lock> ExportableRequirements<'lock> {
                         );
                     }
                 }
+                if !root_extras.is_empty() {
+                    selected_extras.insert(&dist.id, root_extras);
+                }
             }
 
             // Add any development dependencies.
@@ -271,13 +278,49 @@ impl<'lock> ExportableRequirements<'lock> {
                     .or_insert_with(|| graph.add_node(Node::Package(dep_dist)));
 
                 // Add an edge from the dependency.
+                let mut active_extras = selected_extras
+                    .get(&package.id)
+                    .cloned()
+                    .unwrap_or_default();
+                if let Some(extra) = extra {
+                    if !active_extras.contains(extra) {
+                        active_extras.push((*extra).clone());
+                    }
+                }
+                if !active_extras.is_empty() {
+                    match selected_extras.entry(&dep.package_id) {
+                        Entry::Occupied(mut entry) => {
+                            for extra in &active_extras {
+                                if !entry.get().contains(extra) {
+                                    entry.get_mut().push(extra.clone());
+                                }
+                            }
+                        }
+                        Entry::Vacant(entry) => {
+                            entry.insert(active_extras.clone());
+                        }
+                    }
+                }
+
                 graph.add_edge(
                     index,
                     dep_index,
                     if let Some(extra) = extra {
-                        Edge::Optional(extra, dep.simplified_marker.as_simplified_marker_tree())
+                        Edge::Optional(
+                            extra,
+                            dep.simplified_marker
+                                .as_simplified_marker_tree()
+                                .simplify_extras(std::slice::from_ref(extra)),
+                        )
                     } else {
-                        Edge::Prod(dep.simplified_marker.as_simplified_marker_tree())
+                        Edge::Prod(selected_extras.get(&package.id).map_or_else(
+                            || dep.simplified_marker.as_simplified_marker_tree(),
+                            |extras| {
+                                dep.simplified_marker
+                                    .as_simplified_marker_tree()
+                                    .simplify_extras(extras)
+                            },
+                        ))
                     },
                 );
 

crates/uv/tests/it/export.rs

@@ -344,6 +344,81 @@ fn requirements_txt_project_extra() -> Result<()> {
     Ok(())
 }
 
+#[test]
+fn requirements_txt_simplifies_selected_root_extra_markers_from_lock() -> Result<()> {
+    let context = uv_test::test_context!("3.12");
+
+    let pyproject_toml = context.temp_dir.child("pyproject.toml");
+    pyproject_toml.write_str(
+        r#"
+        [project]
+        name = "project"
+        version = "0.1.0"
+        requires-python = ">=3.12"
+        dependencies = ["typing-extensions ; extra == 'foo' or extra == 'bar'"]
+
+        [project.optional-dependencies]
+        foo = []
+        bar = []
+        "#,
+    )?;
+
+    let lock = context.temp_dir.child("uv.lock");
+    lock.write_str(
+        r#"
+        version = 1
+        revision = 3
+        requires-python = ">=3.12"
+
+        [[package]]
+        name = "project"
+        version = "0.1.0"
+        source = { virtual = "." }
+        dependencies = [
+            { name = "typing-extensions", marker = "extra == 'foo' or extra == 'bar'" },
+        ]
+
+        [package.optional-dependencies]
+        foo = []
+        bar = []
+
+        [package.metadata]
+        requires-dist = [
+            { name = "typing-extensions", marker = "extra == 'foo' or extra == 'bar'" },
+        ]
+        provides-extras = ["foo", "bar"]
+
+        [[package]]
+        name = "typing-extensions"
+        version = "4.10.0"
+        source = { registry = "https://pypi.org/simple" }
+        dependencies = []
+        "#,
+    )?;
+
+    uv_snapshot!(
+        context.filters(),
+        context
+            .export()
+            .arg("--all-extras")
+            .arg("--frozen")
+            .arg("--no-hashes"),
+        @r"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+    # This file was autogenerated by uv via the following command:
+    #    uv export --cache-dir [CACHE_DIR] --all-extras --frozen --no-hashes
+    typing-extensions==4.10.0
+        # via project
+
+    ----- stderr -----
+    "
+    );
+
+    Ok(())
+}
+
 #[test]
 fn requirements_txt_prune() -> Result<()> {
     let context = uv_test::test_context!("3.12");