Include newly-added optional dependencies in lockfile (#5686)

## Summary

When we add a new optional group in `uv add`, we never to update the
`pyproject.toml` before locking. Otherwise, we use the stale
`pyproject.toml` and omit the optional group.

Closes https://github.com/astral-sh/uv/issues/5687.

crates/uv-workspace/src/pyproject_mut.rs

@@ -24,6 +24,8 @@ pub enum Error {
     Parse(#[from] Box<TomlError>),
     #[error("Failed to serialize `pyproject.toml`")]
     Serialize(#[from] Box<toml::ser::Error>),
+    #[error("Failed to deserialize `pyproject.toml`")]
+    Deserialize(#[from] Box<toml::de::Error>),
     #[error("Dependencies in `pyproject.toml` are malformed")]
     MalformedDependencies,
     #[error("Sources in `pyproject.toml` are malformed")]
@@ -35,13 +37,18 @@ pub enum Error {
 }
 
 impl PyProjectTomlMut {
-    /// Initialize a `PyProjectTomlMut` from a `PyProjectToml`.
+    /// Initialize a [`PyProjectTomlMut`] from a [`PyProjectToml`].
     pub fn from_toml(pyproject: &PyProjectToml) -> Result<Self, Error> {
         Ok(Self {
             doc: pyproject.raw.parse().map_err(Box::new)?,
         })
     }
 
+    /// Initialize a [`PyProjectToml`] from a [`PyProjectTomlMut`].
+    pub fn to_toml(&self) -> Result<PyProjectToml, Error> {
+        Ok(toml::from_str(&self.doc.to_string()).map_err(Box::new)?)
+    }
+
     /// Adds a project to the workspace.
     pub fn add_workspace(&mut self, path: impl AsRef<Path>) -> Result<(), Error> {
         // Get or create `tool.uv.workspace.members`.

crates/uv-workspace/src/workspace.rs

@@ -197,6 +197,48 @@ impl Workspace {
         })
     }
 
+    /// Set the [`ProjectWorkspace`] for a given workspace member.
+    ///
+    /// Assumes that the project name is unchanged in the updated [`PyProjectToml`].
+    #[must_use]
+    pub fn with_pyproject_toml(
+        self,
+        package_name: &PackageName,
+        pyproject_toml: PyProjectToml,
+    ) -> Option<Self> {
+        let mut packages = self.packages;
+        let member = packages.get_mut(package_name)?;
+
+        if member.root == self.install_path {
+            // If the member is also the workspace root, update _both_ the member entry and the
+            // root `pyproject.toml`.
+            let workspace_pyproject_toml = pyproject_toml.clone();
+
+            // Refresh the workspace sources.
+            let workspace_sources = workspace_pyproject_toml
+                .tool
+                .clone()
+                .and_then(|tool| tool.uv)
+                .and_then(|uv| uv.sources)
+                .unwrap_or_default();
+
+            // Set the `pyproject.toml` for the member.
+            member.pyproject_toml = pyproject_toml;
+
+            Some(Self {
+                pyproject_toml: workspace_pyproject_toml,
+                sources: workspace_sources,
+                packages,
+                ..self
+            })
+        } else {
+            // Set the `pyproject.toml` for the member.
+            member.pyproject_toml = pyproject_toml;
+
+            Some(Self { packages, ..self })
+        }
+    }
+
     /// Returns the set of requirements that include all packages in the workspace.
     pub fn members_as_requirements(&self) -> Vec<Requirement> {
         self.packages
@@ -765,6 +807,19 @@ impl ProjectWorkspace {
         &self.workspace().packages[&self.project_name]
     }
 
+    /// Set the `pyproject.toml` for the current project.
+    ///
+    /// Assumes that the project name is unchanged in the updated [`PyProjectToml`].
+    #[must_use]
+    pub fn with_pyproject_toml(self, pyproject_toml: PyProjectToml) -> Option<Self> {
+        Some(Self {
+            workspace: self
+                .workspace
+                .with_pyproject_toml(&self.project_name, pyproject_toml)?,
+            ..self
+        })
+    }
+
     /// Find the workspace for a project.
     pub async fn from_project(
         install_path: &Path,

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

@@ -219,6 +219,12 @@ pub(crate) async fn add(
         return Ok(ExitStatus::Success);
     }
 
+    // Update the `pypackage.toml` in-memory.
+    let project = project
+        .clone()
+        .with_pyproject_toml(pyproject.to_toml()?)
+        .context("Failed to update `pyproject.toml`")?;
+
     // Initialize any shared state.
     let state = SharedState::default();
 

crates/uv/tests/edit.rs

@@ -834,10 +834,13 @@ fn add_remove_optional() -> Result<()> {
 
     ----- stderr -----
     warning: `uv add` is experimental and may change without warning
-    Resolved 1 package in [TIME]
-    Prepared 1 package in [TIME]
-    Installed 1 package in [TIME]
+    Resolved 4 packages in [TIME]
+    Prepared 4 packages in [TIME]
+    Installed 4 packages in [TIME]
+     + anyio==3.7.0
+     + idna==3.6
      + project==0.1.0 (from file://[TEMP_DIR]/)
+     + sniffio==1.3.1
     "###);
 
     let pyproject_toml = fs_err::read_to_string(context.temp_dir.join("pyproject.toml"))?;
@@ -873,15 +876,52 @@ fn add_remove_optional() -> Result<()> {
         requires-python = ">=3.12"
         exclude-newer = "2024-03-25 00:00:00 UTC"
 
+        [[distribution]]
+        name = "anyio"
+        version = "3.7.0"
+        source = { registry = "https://pypi.org/simple" }
+        dependencies = [
+            { name = "idna" },
+            { name = "sniffio" },
+        ]
+        sdist = { url = "https://files.pythonhosted.org/packages/c6/b3/fefbf7e78ab3b805dec67d698dc18dd505af7a18a8dd08868c9b4fa736b5/anyio-3.7.0.tar.gz", hash = "sha256:275d9973793619a5374e1c89a4f4ad3f4b0a5510a2b5b939444bee8f4c4d37ce", size = 142737 }
+        wheels = [
+            { url = "https://files.pythonhosted.org/packages/68/fe/7ce1926952c8a403b35029e194555558514b365ad77d75125f521a2bec62/anyio-3.7.0-py3-none-any.whl", hash = "sha256:eddca883c4175f14df8aedce21054bfca3adb70ffe76a9f607aef9d7fa2ea7f0", size = 80873 },
+        ]
+
+        [[distribution]]
+        name = "idna"
+        version = "3.6"
+        source = { registry = "https://pypi.org/simple" }
+        sdist = { url = "https://files.pythonhosted.org/packages/bf/3f/ea4b9117521a1e9c50344b909be7886dd00a519552724809bb1f486986c2/idna-3.6.tar.gz", hash = "sha256:9ecdbbd083b06798ae1e86adcbfe8ab1479cf864e4ee30fe4e46a003d12491ca", size = 175426 }
+        wheels = [
+            { url = "https://files.pythonhosted.org/packages/c2/e7/a82b05cf63a603df6e68d59ae6a68bf5064484a0718ea5033660af4b54a9/idna-3.6-py3-none-any.whl", hash = "sha256:c05567e9c24a6b9faaa835c4821bad0590fbb9d5779e7caa6e1cc4978e7eb24f", size = 61567 },
+        ]
+
         [[distribution]]
         name = "project"
         version = "0.1.0"
         source = { editable = "." }
+
+        [distribution.optional-dependencies]
+        io = [
+            { name = "anyio" },
+        ]
+
+        [[distribution]]
+        name = "sniffio"
+        version = "1.3.1"
+        source = { registry = "https://pypi.org/simple" }
+        sdist = { url = "https://files.pythonhosted.org/packages/a2/87/a6771e1546d97e7e041b6ae58d80074f81b7d5121207425c964ddf5cfdbd/sniffio-1.3.1.tar.gz", hash = "sha256:f4324edc670a0f49750a81b895f35c3adb843cca46f0530f79fc1babb23789dc", size = 20372 }
+        wheels = [
+            { url = "https://files.pythonhosted.org/packages/e9/44/75a9c9421471a6c4805dbf2356f7c181a29c1879239abab1ea2cc8f38b40/sniffio-1.3.1-py3-none-any.whl", hash = "sha256:2f6da418d1f1e0fddd844478f41680e794e6051915791a034ff65e5f100525a2", size = 10235 },
+        ]
         "###
         );
     });
 
-    // Install from the lockfile.
+    // Install from the lockfile. At present, this will _uninstall_ the packages since `sync` does
+    // not include extras by default.
     uv_snapshot!(context.filters(), context.sync().arg("--frozen"), @r###"
     success: true
     exit_code: 0
@@ -889,7 +929,10 @@ fn add_remove_optional() -> Result<()> {
 
     ----- stderr -----
     warning: `uv sync` is experimental and may change without warning
-    Audited 1 package in [TIME]
+    Uninstalled 3 packages in [TIME]
+     - anyio==3.7.0
+     - idna==3.6
+     - sniffio==1.3.1
     "###);
 
     // This should fail without --optional.