Separate requirements.txt export logic from graph construction

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

@@ -1,52 +1,50 @@
-use std::borrow::Cow;
 use std::collections::hash_map::Entry;
 use std::collections::VecDeque;
-use std::fmt::Formatter;
-use std::path::{Component, Path, PathBuf};
 
 use either::Either;
-use owo_colors::OwoColorize;
 use petgraph::graph::NodeIndex;
 use petgraph::prelude::EdgeRef;
 use petgraph::visit::IntoNodeReferences;
 use petgraph::{Direction, Graph};
 use rustc_hash::{FxBuildHasher, FxHashMap, FxHashSet};
-use url::Url;
-
-use uv_configuration::{
-    DependencyGroupsWithDefaults, EditableMode, ExtrasSpecification, InstallOptions,
-};
-use uv_distribution_filename::{DistExtension, SourceDistExtension};
-use uv_fs::Simplified;
-use uv_git_types::GitReference;
+
+use uv_configuration::{DependencyGroupsWithDefaults, ExtrasSpecification, InstallOptions};
 use uv_normalize::{ExtraName, GroupName, PackageName};
 use uv_pep508::MarkerTree;
-use uv_pypi_types::{ConflictItem, ParsedArchiveUrl, ParsedGitUrl};
+use uv_pypi_types::ConflictItem;
 
 use crate::graph_ops::{marker_reachability, Reachable};
-use crate::lock::{Package, PackageId, Source};
+pub use crate::lock::export::requirements_txt::RequirementsTxtExport;
 use crate::universal_marker::resolve_conflicts;
-use crate::{Installable, LockError};
-
-/// An export of a [`Lock`] that renders in `requirements.txt` format.
-#[derive(Debug)]
-pub struct RequirementsTxtExport<'lock> {
-    nodes: Vec<Requirement<'lock>>,
-    hashes: bool,
-    editable: EditableMode,
+use crate::{Installable, Package};
+
+mod requirements_txt;
+
+/// A flat requirement, with its associated marker.
+#[derive(Debug, Clone, PartialEq, Eq)]
+struct ExportableRequirement<'lock> {
+    /// The [`Package`] associated with the requirement.
+    package: &'lock Package,
+    /// The marker that must be satisfied to install the package.
+    marker: MarkerTree,
+    /// The list of packages that depend on this package.
+    dependents: Vec<&'lock Package>,
 }
 
-impl<'lock> RequirementsTxtExport<'lock> {
-    pub fn from_lock(
+/// A set of flattened, exportable requirements, generated from a lockfile.
+#[derive(Debug, Clone, PartialEq, Eq)]
+struct ExportableRequirements<'lock>(Vec<ExportableRequirement<'lock>>);
+
+impl<'lock> ExportableRequirements<'lock> {
+    /// Generate the set of exportable [`ExportableRequirement`] entries from the given lockfile.
+    fn from_lock(
         target: &impl Installable<'lock>,
         prune: &[PackageName],
         extras: &ExtrasSpecification,
         dev: &DependencyGroupsWithDefaults,
         annotate: bool,
-        editable: EditableMode,
-        hashes: bool,
         install_options: &'lock InstallOptions,
-    ) -> Result<Self, LockError> {
+    ) -> Self {
         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);
@@ -292,7 +290,7 @@ impl<'lock> RequirementsTxtExport<'lock> {
         };
 
         // Collect all packages.
-        let mut nodes = graph
+        let nodes = graph
             .node_references()
             .filter_map(|(index, node)| match node {
                 Node::Root => None,
@@ -305,7 +303,7 @@ impl<'lock> RequirementsTxtExport<'lock> {
                     target.lock().members(),
                 )
             })
-            .map(|(index, package)| Requirement {
+            .map(|(index, package)| ExportableRequirement {
                 package,
                 marker: reachability.remove(&index).unwrap_or_default(),
                 dependents: if annotate {
@@ -327,16 +325,47 @@ impl<'lock> RequirementsTxtExport<'lock> {
             .filter(|requirement| !requirement.marker.is_false())
             .collect::<Vec<_>>();
 
-        // Sort the nodes, such that unnamed URLs (editables) appear at the top.
-        nodes.sort_unstable_by(|a, b| {
-            RequirementComparator::from(a.package).cmp(&RequirementComparator::from(b.package))
-        });
+        Self(nodes)
+    }
+}
 
-        Ok(Self {
-            nodes,
-            hashes,
-            editable,
-        })
+/// A node in the graph.
+#[derive(Debug, Clone, PartialEq, Eq)]
+enum Node<'lock> {
+    Root,
+    Package(&'lock Package),
+}
+
+/// An edge in the resolution graph, along with the marker that must be satisfied to traverse it.
+#[derive(Debug, Clone)]
+enum Edge<'lock> {
+    Prod(MarkerTree),
+    Optional(&'lock ExtraName, MarkerTree),
+    Dev(&'lock GroupName, MarkerTree),
+}
+
+impl Edge<'_> {
+    /// Return the [`MarkerTree`] for this edge.
+    fn marker(&self) -> &MarkerTree {
+        match self {
+            Self::Prod(marker) => marker,
+            Self::Optional(_, marker) => marker,
+            Self::Dev(_, marker) => marker,
+        }
+    }
+}
+
+impl Reachable<MarkerTree> for Edge<'_> {
+    fn true_marker() -> MarkerTree {
+        MarkerTree::TRUE
+    }
+
+    fn false_marker() -> MarkerTree {
+        MarkerTree::FALSE
+    }
+
+    fn marker(&self) -> MarkerTree {
+        *self.marker()
     }
 }
 
@@ -502,197 +531,3 @@ fn conflict_marker_reachability<'lock>(
 
     reachability
 }
-
-impl std::fmt::Display for RequirementsTxtExport<'_> {
-    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        // Write out each package.
-        for Requirement {
-            package,
-            marker,
-            dependents,
-        } in &self.nodes
-        {
-            match &package.id.source {
-                Source::Registry(_) => {
-                    let version = package
-                        .id
-                        .version
-                        .as_ref()
-                        .expect("registry package without version");
-                    write!(f, "{}=={}", package.id.name, version)?;
-                }
-                Source::Git(url, git) => {
-                    // Remove the fragment and query from the URL; they're already present in the
-                    // `GitSource`.
-                    let mut url = url.to_url().map_err(|_| std::fmt::Error)?;
-                    url.set_fragment(None);
-                    url.set_query(None);
-
-                    // Reconstruct the `GitUrl` from the `GitSource`.
-                    let git_url = uv_git_types::GitUrl::from_commit(
-                        url,
-                        GitReference::from(git.kind.clone()),
-                        git.precise,
-                    )
-                    .expect("Internal Git URLs must have supported schemes");
-
-                    // Reconstruct the PEP 508-compatible URL from the `GitSource`.
-                    let url = Url::from(ParsedGitUrl {
-                        url: git_url.clone(),
-                        subdirectory: git.subdirectory.clone(),
-                    });
-
-                    write!(f, "{} @ {}", package.id.name, url)?;
-                }
-                Source::Direct(url, direct) => {
-                    let url = Url::from(ParsedArchiveUrl {
-                        url: url.to_url().map_err(|_| std::fmt::Error)?,
-                        subdirectory: direct.subdirectory.clone(),
-                        ext: DistExtension::Source(SourceDistExtension::TarGz),
-                    });
-                    write!(f, "{} @ {}", package.id.name, url)?;
-                }
-                Source::Path(path) | Source::Directory(path) => {
-                    if path.is_absolute() {
-                        write!(
-                            f,
-                            "{}",
-                            Url::from_file_path(path).map_err(|()| std::fmt::Error)?
-                        )?;
-                    } else {
-                        write!(f, "{}", anchor(path).portable_display())?;
-                    }
-                }
-                Source::Editable(path) => match self.editable {
-                    EditableMode::Editable => {
-                        write!(f, "-e {}", anchor(path).portable_display())?;
-                    }
-                    EditableMode::NonEditable => {
-                        if path.is_absolute() {
-                            write!(
-                                f,
-                                "{}",
-                                Url::from_file_path(path).map_err(|()| std::fmt::Error)?
-                            )?;
-                        } else {
-                            write!(f, "{}", anchor(path).portable_display())?;
-                        }
-                    }
-                },
-                Source::Virtual(_) => {
-                    continue;
-                }
-            }
-
-            if let Some(contents) = marker.contents() {
-                write!(f, " ; {contents}")?;
-            }
-
-            if self.hashes {
-                let mut hashes = package.hashes();
-                hashes.sort_unstable();
-                if !hashes.is_empty() {
-                    for hash in hashes.iter() {
-                        writeln!(f, " \\")?;
-                        write!(f, "    --hash=")?;
-                        write!(f, "{hash}")?;
-                    }
-                }
-            }
-
-            writeln!(f)?;
-
-            // Add "via ..." comments for all dependents.
-            match dependents.as_slice() {
-                [] => {}
-                [dependent] => {
-                    writeln!(f, "{}", format!("    # via {}", dependent.id.name).green())?;
-                }
-                _ => {
-                    writeln!(f, "{}", "    # via".green())?;
-                    for &dependent in dependents {
-                        writeln!(f, "{}", format!("    #   {}", dependent.id.name).green())?;
-                    }
-                }
-            }
-        }
-
-        Ok(())
-    }
-}
-
-/// A node in the graph.
-#[derive(Debug, Clone, PartialEq, Eq)]
-enum Node<'lock> {
-    Root,
-    Package(&'lock Package),
-}
-
-/// An edge in the resolution graph, along with the marker that must be satisfied to traverse it.
-#[derive(Debug, Clone)]
-enum Edge<'lock> {
-    Prod(MarkerTree),
-    Optional(&'lock ExtraName, MarkerTree),
-    Dev(&'lock GroupName, MarkerTree),
-}
-
-impl Edge<'_> {
-    /// Return the [`MarkerTree`] for this edge.
-    fn marker(&self) -> &MarkerTree {
-        match self {
-            Self::Prod(marker) => marker,
-            Self::Optional(_, marker) => marker,
-            Self::Dev(_, marker) => marker,
-        }
-    }
-}
-
-impl Reachable<MarkerTree> for Edge<'_> {
-    fn true_marker() -> MarkerTree {
-        MarkerTree::TRUE
-    }
-
-    fn false_marker() -> MarkerTree {
-        MarkerTree::FALSE
-    }
-
-    fn marker(&self) -> MarkerTree {
-        *self.marker()
-    }
-}
-
-/// A flat requirement, with its associated marker.
-#[derive(Debug, Clone, PartialEq, Eq)]
-struct Requirement<'lock> {
-    package: &'lock Package,
-    marker: MarkerTree,
-    dependents: Vec<&'lock Package>,
-}
-
-#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
-enum RequirementComparator<'lock> {
-    Editable(&'lock Path),
-    Path(&'lock Path),
-    Package(&'lock PackageId),
-}
-
-impl<'lock> From<&'lock Package> for RequirementComparator<'lock> {
-    fn from(value: &'lock Package) -> Self {
-        match &value.id.source {
-            Source::Path(path) | Source::Directory(path) => Self::Path(path),
-            Source::Editable(path) => Self::Editable(path),
-            _ => Self::Package(&value.id),
-        }
-    }
-}
-
-/// Modify a relative [`Path`] to anchor it at the current working directory.
-///
-/// For example, given `foo/bar`, returns `./foo/bar`.
-fn anchor(path: &Path) -> Cow<'_, Path> {
-    match path.components().next() {
-        None => Cow::Owned(PathBuf::from(".")),
-        Some(Component::CurDir | Component::ParentDir) => Cow::Borrowed(path),
-        _ => Cow::Owned(PathBuf::from("./").join(path)),
-    }
-}