From 0f9f37d82ed9910f25116440c5b03097d6d5bd32 Mon Sep 17 00:00:00 2001
From: j178 <10510431+j178@users.noreply.github.com>
Date: Tue, 23 Jul 2024 00:18:58 +0800
Subject: [PATCH] Add to root member

---
 crates/uv-workspace/src/workspace.rs   | 11 +++++++++++
 crates/uv/src/commands/project/init.rs | 12 +++++-------
 2 files changed, 16 insertions(+), 7 deletions(-)

diff --git a/crates/uv-workspace/src/workspace.rs b/crates/uv-workspace/src/workspace.rs
index edeb36a4e490..6f0b31fe48fc 100644
--- a/crates/uv-workspace/src/workspace.rs
+++ b/crates/uv-workspace/src/workspace.rs
@@ -674,6 +674,17 @@ impl ProjectWorkspace {
         &self.workspace().packages[&self.project_name]
     }
 
+    /// Returns the workspace root member as a [`WorkspaceMember`].
+    pub fn root_member(&self) -> &WorkspaceMember {
+        self.workspace().packages.iter().find_map(|(_, member)| {
+            if member.root() == self.workspace().install_path() {
+                Some(member)
+            } else {
+                None
+            }
+        }).expect("workspace root member must exist")
+    }
+
     /// Find the workspace for a project.
     pub async fn from_project(
         install_path: &Path,
diff --git a/crates/uv/src/commands/project/init.rs b/crates/uv/src/commands/project/init.rs
index 76e6e66fe3ab..dd6bb833dd6f 100644
--- a/crates/uv/src/commands/project/init.rs
+++ b/crates/uv/src/commands/project/init.rs
@@ -111,13 +111,13 @@ pub(crate) async fn init(
 
     if let Some(workspace) = workspace {
         // Add the package to the workspace.
-        let mut pyproject =
-            PyProjectTomlMut::from_toml(workspace.current_project().pyproject_toml())?;
-        pyproject.add_workspace(path.strip_prefix(workspace.project_root())?)?;
+        let root_member = workspace.root_member();
+        let mut pyproject = PyProjectTomlMut::from_toml(root_member.pyproject_toml())?;
+        pyproject.add_workspace(path.strip_prefix(root_member.root())?)?;
 
         // Save the modified `pyproject.toml`.
         fs_err::write(
-            workspace.current_project().root().join("pyproject.toml"),
+            root_member.root().join("pyproject.toml"),
             pyproject.to_string(),
         )?;
 
@@ -125,9 +125,7 @@ pub(crate) async fn init(
             printer.stderr(),
             "Adding {} as member of workspace {}",
             name.cyan(),
-            workspace
-                .workspace()
-                .install_path()
+            root_member.root()
                 .simplified_display()
                 .cyan()
         )?;