Support PEP 723 scripts in uv add and uv remove

crates/uv-cli/src/lib.rs

@@ -2462,6 +2462,10 @@ pub struct AddArgs {
         help_heading = "Python options"
     )]
     pub python: Option<String>,
+
+    /// Specifies the Python script where the dependency will be added.
+    #[arg(long)]
+    pub script: Option<PathBuf>,
 }
 
 #[derive(Args)]
@@ -2509,6 +2513,9 @@ pub struct RemoveArgs {
     #[arg(long, conflicts_with = "isolated")]
     pub package: Option<PackageName>,
 
+    /// Specifies the Python script where the dependency will be removed.
+    #[arg(long)]
+    pub script: Option<PathBuf>,
     /// The Python interpreter to use for resolving and syncing.
     ///
     /// See `uv help python` for details on Python discovery and supported

crates/uv-scripts/src/lib.rs

@@ -19,19 +19,43 @@ static FINDER: LazyLock<Finder> = LazyLock::new(|| Finder::new(b"# /// script"))
 pub struct Pep723Script {
     pub path: PathBuf,
     pub metadata: Pep723Metadata,
+    pub data: String,
 }
 
 impl Pep723Script {
     /// Read the PEP 723 `script` metadata from a Python file, if it exists.
     ///
     /// See: <https://peps.python.org/pep-0723/>
     pub async fn read(file: impl AsRef<Path>) -> Result<Option<Self>, Pep723Error> {
-        let metadata = Pep723Metadata::read(&file).await?;
-        Ok(metadata.map(|metadata| Self {
+        let contents = match fs_err::tokio::read(&file).await {
+            Ok(contents) => contents,
+            Err(err) if err.kind() == io::ErrorKind::NotFound => return Ok(None),
+            Err(err) => return Err(err.into()),
+        };
+
+        // Extract the `script` tag.
+        let Some((metadata, data)) = extract_script_tag(&contents)? else {
+            return Ok(None);
+        };
+
+        // Parse the metadata.
+        let metadata = Pep723Metadata::from_string(metadata)?;
+
+        Ok(Some(Self {
             path: file.as_ref().to_path_buf(),
             metadata,
+            data,
         }))
     }
+
+    /// Replace the existing metadata in the file with new metadata and write the updated content.
+    pub async fn replace_metadata(&self, new_metadata: &str) -> Result<(), Pep723Error> {
+        let new_content = format!("{}{}", serialize_metadata(new_metadata), self.data);
+
+        fs_err::tokio::write(&self.path, new_content)
+            .await
+            .map_err(std::convert::Into::into)
+    }
 }
 
 /// PEP 723 metadata as parsed from a `script` comment block.
@@ -43,28 +67,16 @@ pub struct Pep723Metadata {
     pub dependencies: Option<Vec<pep508_rs::Requirement<VerbatimParsedUrl>>>,
     pub requires_python: Option<pep440_rs::VersionSpecifiers>,
     pub tool: Option<Tool>,
+    /// The raw unserialized document.
+    #[serde(skip)]
+    pub raw: String,
 }
 
 impl Pep723Metadata {
-    /// Read the PEP 723 `script` metadata from a Python file, if it exists.
-    ///
-    /// See: <https://peps.python.org/pep-0723/>
-    pub async fn read(file: impl AsRef<Path>) -> Result<Option<Self>, Pep723Error> {
-        let contents = match fs_err::tokio::read(file).await {
-            Ok(contents) => contents,
-            Err(err) if err.kind() == io::ErrorKind::NotFound => return Ok(None),
-            Err(err) => return Err(err.into()),
-        };
-
-        // Extract the `script` tag.
-        let Some(contents) = extract_script_tag(&contents)? else {
-            return Ok(None);
-        };
-
-        // Parse the metadata.
-        let metadata = toml::from_str(&contents)?;
-
-        Ok(Some(metadata))
+    /// Parse `Pep723Metadata` from a raw TOML string.
+    pub fn from_string(raw: String) -> Result<Self, toml::de::Error> {
+        let metadata = toml::from_str(&raw)?;
+        Ok(Pep723Metadata { raw, ..metadata })
     }
 }
 
@@ -94,34 +106,11 @@ pub enum Pep723Error {
     Toml(#[from] toml::de::Error),
 }
 
-/// Read the PEP 723 `script` metadata from a Python file, if it exists.
-///
-/// See: <https://peps.python.org/pep-0723/>
-pub async fn read_pep723_metadata(
-    file: impl AsRef<Path>,
-) -> Result<Option<Pep723Metadata>, Pep723Error> {
-    let contents = match fs_err::tokio::read(file).await {
-        Ok(contents) => contents,
-        Err(err) if err.kind() == io::ErrorKind::NotFound => return Ok(None),
-        Err(err) => return Err(err.into()),
-    };
-
-    // Extract the `script` tag.
-    let Some(contents) = extract_script_tag(&contents)? else {
-        return Ok(None);
-    };
-
-    // Parse the metadata.
-    let metadata = toml::from_str(&contents)?;
-
-    Ok(Some(metadata))
-}
-
 /// Given the contents of a Python file, extract the `script` metadata block, with leading comment
-/// hashes removed.
+/// hashes removed and the python script.
 ///
 /// See: <https://peps.python.org/pep-0723/>
-fn extract_script_tag(contents: &[u8]) -> Result<Option<String>, Pep723Error> {
+fn extract_script_tag(contents: &[u8]) -> Result<Option<(String, String)>, Pep723Error> {
     // Identify the opening pragma.
     let Some(index) = FINDER.find(contents) else {
         return Ok(None);
@@ -149,9 +138,14 @@ fn extract_script_tag(contents: &[u8]) -> Result<Option<String>, Pep723Error> {
     // > second character is a space, otherwise just the first character (which means the line
     // > consists of only a single #).
     let mut toml = vec![];
-    for line in lines {
+
+    let mut python_script = vec![];
+
+    while let Some(line) = lines.next() {
         // Remove the leading `#`.
         let Some(line) = line.strip_prefix('#') else {
+            python_script.push(line);
+            python_script.extend(lines);
             break;
         };
 
@@ -163,11 +157,13 @@ fn extract_script_tag(contents: &[u8]) -> Result<Option<String>, Pep723Error> {
 
         // Otherwise, the line _must_ start with ` `.
         let Some(line) = line.strip_prefix(' ') else {
+            python_script.push(line);
+            python_script.extend(lines);
             break;
         };
+
         toml.push(line);
     }
-
     // Find the closing `# ///`. The precedence is such that we need to identify the _last_ such
     // line.
     //
@@ -202,12 +198,36 @@ fn extract_script_tag(contents: &[u8]) -> Result<Option<String>, Pep723Error> {
 
     // Join the lines into a single string.
     let toml = toml.join("\n") + "\n";
+    let python_script = python_script.join("\n") + "\n";
+
+    Ok(Some((toml, python_script)))
+}
+
+/// Formats the provided metadata by prefixing each line with `#` and wrapping it with script markers.
+fn serialize_metadata(metadata: &str) -> String {
+    let mut output = String::with_capacity(metadata.len() + 2);
+
+    output.push_str("# /// script\n");
 
-    Ok(Some(toml))
+    for line in metadata.lines() {
+        if line.is_empty() {
+            output.push('\n');
+        } else {
+            output.push_str("# ");
+            output.push_str(line);
+            output.push('\n');
+        }
+    }
+
+    output.push_str("# ///\n");
+
+    output
 }
 
 #[cfg(test)]
 mod tests {
+    use crate::serialize_metadata;
+
     #[test]
     fn missing_space() {
         let contents = indoc::indoc! {r"
@@ -269,21 +289,37 @@ mod tests {
             #   'rich',
             # ]
             # ///
+
+            import requests
+            from rich.pretty import pprint
+
+            resp = requests.get('https://peps.python.org/api/peps.json')
+            data = resp.json()
         "};
 
-        let expected = indoc::indoc! {r"
+        let expected_metadata = indoc::indoc! {r"
             requires-python = '>=3.11'
             dependencies = [
               'requests<3',
               'rich',
             ]
         "};
 
+        let expected_data = indoc::indoc! {r"
+
+            import requests
+            from rich.pretty import pprint
+
+            resp = requests.get('https://peps.python.org/api/peps.json')
+            data = resp.json()
+        "};
+
         let actual = super::extract_script_tag(contents.as_bytes())
             .unwrap()
             .unwrap();
 
-        assert_eq!(actual, expected);
+        assert_eq!(actual.0, expected_metadata);
+        assert_eq!(actual.1, expected_data);
     }
 
     #[test]
@@ -312,7 +348,8 @@ mod tests {
 
         let actual = super::extract_script_tag(contents.as_bytes())
             .unwrap()
-            .unwrap();
+            .unwrap()
+            .0;
 
         assert_eq!(actual, expected);
     }
@@ -341,8 +378,42 @@ mod tests {
 
         let actual = super::extract_script_tag(contents.as_bytes())
             .unwrap()
-            .unwrap();
+            .unwrap()
+            .0;
 
         assert_eq!(actual, expected);
     }
+
+    #[test]
+    fn test_serialize_metadata_formatting() {
+        let metadata = indoc::indoc! {r"
+            requires-python = '>=3.11'
+            dependencies = [
+              'requests<3',
+              'rich',
+            ]
+        "};
+
+        let expected_output = indoc::indoc! {r"
+            # /// script
+            # requires-python = '>=3.11'
+            # dependencies = [
+            #   'requests<3',
+            #   'rich',
+            # ]
+            # ///
+        "};
+
+        let result = serialize_metadata(metadata);
+        assert_eq!(result, expected_output);
+    }
+
+    #[test]
+    fn test_serialize_metadata_empty() {
+        let metadata = "";
+        let expected_output = "# /// script\n# ///\n";
+
+        let result = serialize_metadata(metadata);
+        assert_eq!(result, expected_output);
+    }
 }