Include --branch et al when resolving unnamed URLs in uv add (#7447)

## Summary

Closes #7433.

crates/uv-git/src/lib.rs

@@ -60,6 +60,13 @@ impl GitUrl {
         self
     }
 
+    /// Set the [`GitReference`] to use for this Git URL.
+    #[must_use]
+    pub fn with_reference(mut self, reference: GitReference) -> Self {
+        self.reference = reference;
+        self
+    }
+
     /// Return the [`Url`] of the Git repository.
     pub fn repository(&self) -> &Url {
         &self.repository

crates/uv/src/commands/project/add.rs

@@ -10,8 +10,8 @@ use tracing::debug;
 
 use cache_key::RepositoryUrl;
 use distribution_types::UnresolvedRequirement;
-use pep508_rs::{ExtraName, Requirement, VersionOrUrl};
+use pep508_rs::{ExtraName, Requirement, UnnamedRequirement, VersionOrUrl};
-use pypi_types::redact_git_credentials;
+use pypi_types::{redact_git_credentials, ParsedUrl, RequirementSource, VerbatimParsedUrl};
 use uv_auth::{store_credentials_from_url, Credentials};
 use uv_cache::Cache;
 use uv_client::{BaseClientBuilder, Connectivity, FlatIndexClient, RegistryClientBuilder};
@@ -21,7 +21,7 @@ use uv_configuration::{
 use uv_dispatch::BuildDispatch;
 use uv_distribution::DistributionDatabase;
 use uv_fs::{Simplified, CWD};
-use uv_git::GIT_STORE;
+use uv_git::{GitReference, GIT_STORE};
 use uv_normalize::PackageName;
 use uv_python::{
     EnvironmentPreference, Interpreter, PythonDownloads, PythonEnvironment, PythonInstallation,
@@ -317,17 +317,22 @@ pub(crate) async fn add(
     // Resolve any unnamed requirements.
     let requirements = {
         // Partition the requirements into named and unnamed requirements.
-        let (mut requirements, unnamed): (Vec<_>, Vec<_>) =
+        let (mut requirements, unnamed): (Vec<_>, Vec<_>) = requirements
-            requirements
+            .into_iter()
-                .into_iter()
+            .map(|spec| {
-                .partition_map(|spec| match spec.requirement {
+                augment_requirement(
-                    UnresolvedRequirement::Named(requirement) => {
+                    spec.requirement,
-                        itertools::Either::Left(requirement)
+                    rev.as_deref(),
-                    }
+                    tag.as_deref(),
-                    UnresolvedRequirement::Unnamed(requirement) => {
+                    branch.as_deref(),
-                        itertools::Either::Right(requirement)
+                )
-                    }
+            })
-                });
+            .partition_map(|requirement| match requirement {
+                UnresolvedRequirement::Named(requirement) => itertools::Either::Left(requirement),
+                UnresolvedRequirement::Unnamed(requirement) => {
+                    itertools::Either::Right(requirement)
+                }
+            });
 
         // Resolve any unnamed requirements.
         if !unnamed.is_empty() {
@@ -766,6 +771,74 @@ async fn lock_and_sync(
     Ok(())
 }
 
+/// Augment a user-provided requirement by attaching any specification data that was provided
+/// separately from the requirement itself (e.g., `--branch main`).
+fn augment_requirement(
+    requirement: UnresolvedRequirement,
+    rev: Option<&str>,
+    tag: Option<&str>,
+    branch: Option<&str>,
+) -> UnresolvedRequirement {
+    match requirement {
+        UnresolvedRequirement::Named(requirement) => {
+            UnresolvedRequirement::Named(pypi_types::Requirement {
+                source: match requirement.source {
+                    RequirementSource::Git {
+                        repository,
+                        reference,
+                        precise,
+                        subdirectory,
+                        url,
+                    } => {
+                        let reference = if let Some(rev) = rev {
+                            GitReference::from_rev(rev.to_string())
+                        } else if let Some(tag) = tag {
+                            GitReference::Tag(tag.to_string())
+                        } else if let Some(branch) = branch {
+                            GitReference::Branch(branch.to_string())
+                        } else {
+                            reference
+                        };
+                        RequirementSource::Git {
+                            repository,
+                            reference,
+                            precise,
+                            subdirectory,
+                            url,
+                        }
+                    }
+                    _ => requirement.source,
+                },
+                ..requirement
+            })
+        }
+        UnresolvedRequirement::Unnamed(requirement) => {
+            UnresolvedRequirement::Unnamed(UnnamedRequirement {
+                url: match requirement.url.parsed_url {
+                    ParsedUrl::Git(mut git) => {
+                        let reference = if let Some(rev) = rev {
+                            Some(GitReference::from_rev(rev.to_string()))
+                        } else if let Some(tag) = tag {
+                            Some(GitReference::Tag(tag.to_string()))
+                        } else {
+                            branch.map(|branch| GitReference::Branch(branch.to_string()))
+                        };
+                        if let Some(reference) = reference {
+                            git.url = git.url.with_reference(reference);
+                        }
+                        VerbatimParsedUrl {
+                            parsed_url: ParsedUrl::Git(git),
+                            verbatim: requirement.url.verbatim,
+                        }
+                    }
+                    _ => requirement.url,
+                },
+                ..requirement
+            })
+        }
+    }
+}
+
 /// Resolves the source for a requirement and processes it into a PEP 508 compliant format.
 fn resolve_requirement(
     requirement: pypi_types::Requirement,

crates/uv/tests/edit.rs

@@ -584,6 +584,40 @@ fn add_git_error() -> Result<()> {
     Ok(())
 }
 
+#[test]
+#[cfg(feature = "git")]
+fn add_git_branch() -> Result<()> {
+    let context = TestContext::new("3.12");
+
+    let pyproject_toml = context.temp_dir.child("pyproject.toml");
+    pyproject_toml.write_str(indoc! {r#"
+        [project]
+        name = "project"
+        version = "0.1.0"
+        requires-python = ">=3.12"
+        dependencies = []
+
+        [build-system]
+        requires = ["setuptools>=42"]
+        build-backend = "setuptools.build_meta"
+    "#})?;
+
+    uv_snapshot!(context.filters(), context.add().arg("uv-public-pypackage @ git+https://github.com/astral-test/uv-public-pypackage").arg("--branch").arg("test-branch"), @r###"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+
+    ----- stderr -----
+    Resolved 2 packages in [TIME]
+    Prepared 2 packages in [TIME]
+    Installed 2 packages in [TIME]
+     + project==0.1.0 (from file://[TEMP_DIR]/)
+     + uv-public-pypackage==0.1.0 (from git+https://github.com/astral-test/uv-public-pypackage@0dacfd662c64cb4ceb16e6cf65a157a8b715b979)
+    "###);
+
+    Ok(())
+}
+
 /// Add a Git requirement using the `--raw-sources` API.
 #[test]
 #[cfg(feature = "git")]
@@ -1713,6 +1747,12 @@ fn add_workspace_editable() -> Result<()> {
 
     let workspace = context.temp_dir.child("pyproject.toml");
     workspace.write_str(indoc! {r#"
+        [project]
+        name = "parent"
+        version = "0.1.0"
+        requires-python = ">=3.12"
+        dependencies = []
+
         [tool.uv.workspace]
         members = ["child1", "child2"]
     "#})?;
@@ -1753,7 +1793,7 @@ fn add_workspace_editable() -> Result<()> {
     ----- stdout -----
 
     ----- stderr -----
-    Resolved 2 packages in [TIME]
+    Resolved 3 packages in [TIME]
     Prepared 2 packages in [TIME]
     Installed 2 packages in [TIME]
      + child1==0.1.0 (from file://[TEMP_DIR]/child1)
@@ -1803,6 +1843,7 @@ fn add_workspace_editable() -> Result<()> {
         members = [
             "child1",
             "child2",
+            "parent",
         ]
 
         [[package]]
@@ -1820,6 +1861,11 @@ fn add_workspace_editable() -> Result<()> {
         name = "child2"
         version = "0.1.0"
         source = { editable = "child2" }
+
+        [[package]]
+        name = "parent"
+        version = "0.1.0"
+        source = { virtual = "." }
         "###
         );
     });
@@ -1837,6 +1883,124 @@ fn add_workspace_editable() -> Result<()> {
     Ok(())
 }
 
+/// Add a workspace dependency via its path.
+#[test]
+fn add_workspace_path() -> Result<()> {
+    let context = TestContext::new("3.12");
+
+    let workspace = context.temp_dir.child("pyproject.toml");
+    workspace.write_str(indoc! {r#"
+        [project]
+        name = "parent"
+        version = "0.1.0"
+        requires-python = ">=3.12"
+        dependencies = []
+
+        [tool.uv.workspace]
+        members = ["child"]
+    "#})?;
+
+    let pyproject_toml = context.temp_dir.child("child/pyproject.toml");
+    pyproject_toml.write_str(indoc! {r#"
+        [project]
+        name = "child"
+        version = "0.1.0"
+        requires-python = ">=3.12"
+        dependencies = []
+
+        [build-system]
+        requires = ["setuptools>=42"]
+        build-backend = "setuptools.build_meta"
+    "#})?;
+
+    uv_snapshot!(context.filters(), context.add().arg("./child"), @r###"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+
+    ----- stderr -----
+    Resolved 2 packages in [TIME]
+    Prepared 1 package in [TIME]
+    Installed 1 package in [TIME]
+     + child==0.1.0 (from file://[TEMP_DIR]/child)
+    "###);
+
+    let pyproject_toml = fs_err::read_to_string(context.temp_dir.child("pyproject.toml"))?;
+
+    insta::with_settings!({
+        filters => context.filters(),
+    }, {
+        assert_snapshot!(
+            pyproject_toml, @r###"
+        [project]
+        name = "parent"
+        version = "0.1.0"
+        requires-python = ">=3.12"
+        dependencies = [
+            "child",
+        ]
+
+        [tool.uv.workspace]
+        members = ["child"]
+
+        [tool.uv.sources]
+        child = { workspace = true }
+        "###
+        );
+    });
+
+    // `uv add` implies a full lock and sync, including development dependencies.
+    let lock = fs_err::read_to_string(context.temp_dir.join("uv.lock"))?;
+
+    insta::with_settings!({
+        filters => context.filters(),
+    }, {
+        assert_snapshot!(
+            lock, @r###"
+        version = 1
+        requires-python = ">=3.12"
+
+        [options]
+        exclude-newer = "2024-03-25T00:00:00Z"
+
+        [manifest]
+        members = [
+            "child",
+            "parent",
+        ]
+
+        [[package]]
+        name = "child"
+        version = "0.1.0"
+        source = { editable = "child" }
+
+        [[package]]
+        name = "parent"
+        version = "0.1.0"
+        source = { virtual = "." }
+        dependencies = [
+            { name = "child" },
+        ]
+
+        [package.metadata]
+        requires-dist = [{ name = "child", editable = "child" }]
+        "###
+        );
+    });
+
+    // Install from the lockfile.
+    uv_snapshot!(context.filters(), context.sync().arg("--frozen"), @r###"
+    success: true
+    exit_code: 0
+    ----- stdout -----
+
+    ----- stderr -----
+    Audited 1 package in [TIME]
+    "###);
+
+    Ok(())
+}
+
 /// Add a path dependency.
 #[test]
 fn add_path() -> Result<()> {