feat(toml): Allow adding/removing from cargo scripts

Author: epage

src/cargo/util/toml_mut/dependency.rs

@@ -1252,6 +1252,8 @@ mod tests {
         let mut local = LocalManifest {
             path: crate_root.clone(),
             manifest,
+            embedded: None,
+            raw: toml.to_owned(),
         };
         assert_eq!(local.manifest.to_string(), toml);
         let gctx = GlobalContext::default().unwrap();

src/cargo/util/toml_mut/manifest.rs

@@ -11,6 +11,7 @@ use crate::core::dependency::DepKind;
 use crate::core::{FeatureValue, Features, Workspace};
 use crate::util::closest;
 use crate::util::interning::InternedString;
+use crate::util::toml::{is_embedded, ScriptSource};
 use crate::{CargoResult, GlobalContext};
 
 /// Dependency table to add deps to.
@@ -245,6 +246,10 @@ pub struct LocalManifest {
     pub path: PathBuf,
     /// Manifest contents.
     pub manifest: Manifest,
+    /// The raw, unparsed package file
+    pub raw: String,
+    /// Edit location for an embedded manifest, if relevant
+    pub embedded: Option<Embedded>,
 }
 
 impl Deref for LocalManifest {
@@ -267,18 +272,56 @@ impl LocalManifest {
         if !path.is_absolute() {
             anyhow::bail!("can only edit absolute paths, got {}", path.display());
         }
-        let data = cargo_util::paths::read(&path)?;
+        let raw = cargo_util::paths::read(&path)?;
+        let mut data = raw.clone();
+        let mut embedded = None;
+        if is_embedded(path) {
+            let source = ScriptSource::parse(&data)?;
+            if let Some(frontmatter) = source.frontmatter() {
+                embedded = Some(Embedded::exists(&data, frontmatter));
+                data = frontmatter.to_owned();
+            } else if let Some(shebang) = source.shebang() {
+                embedded = Some(Embedded::after(&data, shebang));
+                data = String::new();
+            } else {
+                embedded = Some(Embedded::start());
+                data = String::new();
+            }
+        }
         let manifest = data.parse().context("Unable to parse Cargo.toml")?;
         Ok(LocalManifest {
             manifest,
             path: path.to_owned(),
+            raw,
+            embedded,
         })
     }
 
     /// Write changes back to the file.
     pub fn write(&self) -> CargoResult<()> {
-        let s = self.manifest.data.to_string();
-        let new_contents_bytes = s.as_bytes();
+        let mut manifest = self.manifest.data.to_string();
+        let raw = match self.embedded.as_ref() {
+            Some(Embedded::Implicit(start)) => {
+                if !manifest.ends_with("\n") {
+                    manifest.push_str("\n");
+                }
+                let fence = "---\n";
+                let prefix = &self.raw[0..*start];
+                let suffix = &self.raw[*start..];
+                let empty_line = if prefix.is_empty() { "\n" } else { "" };
+                format!("{prefix}{fence}{manifest}{fence}{empty_line}{suffix}")
+            }
+            Some(Embedded::Explicit(span)) => {
+                if !manifest.ends_with("\n") {
+                    manifest.push_str("\n");
+                }
+                let prefix = &self.raw[0..span.start];
+                let suffix = &self.raw[span.end..];
+                format!("{prefix}{manifest}{suffix}")
+            }
+            None => manifest,
+        };
+        let new_contents_bytes = raw.as_bytes();
 
         cargo_util::paths::write_atomic(&self.path, new_contents_bytes)
     }
@@ -531,6 +574,51 @@ impl std::fmt::Display for LocalManifest {
     }
 }
 
+/// Edit location for an embedded manifest
+#[derive(Clone, Debug)]
+pub enum Embedded {
+    /// Manifest is implicit
+    ///
+    /// This is the insert location for a frontmatter
+    Implicit(usize),
+    /// Manifest is explicit in a frontmatter
+    ///
+    /// This is the span of the frontmatter body
+    Explicit(std::ops::Range<usize>),
+}
+
+impl Embedded {
+    fn start() -> Self {
+        Self::Implicit(0)
+    }
+
+    fn after(input: &str, after: &str) -> Self {
+        let span = substr_span(input, after);
+        let end = span.end;
+        Self::Implicit(end)
+    }
+
+    fn exists(input: &str, exists: &str) -> Self {
+        let span = substr_span(input, exists);
+        Self::Explicit(span)
+    }
+}
+
+fn substr_span(haystack: &str, needle: &str) -> std::ops::Range<usize> {
+    let haystack_start_ptr = haystack.as_ptr();
+    let haystack_end_ptr = haystack[haystack.len()..haystack.len()].as_ptr();
+
+    let needle_start_ptr = needle.as_ptr();
+    let needle_end_ptr = needle[needle.len()..needle.len()].as_ptr();
+
+    assert!(needle_end_ptr < haystack_end_ptr);
+    assert!(haystack_start_ptr <= needle_start_ptr);
+    let start = needle_start_ptr as usize - haystack_start_ptr as usize;
+    let end = start + needle.len();
+
+    start..end
+}
+
 #[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
 enum DependencyStatus {
     None,

tests/testsuite/cargo_add/script_bare/mod.rs

@@ -31,7 +31,7 @@ fn case() {
         .arg_line("--manifest-path cargo-test-fixture.rs my-package")
         .current_dir(cwd)
         .assert()
-        .failure()
+        .success()
         .stdout_eq(str![""])
         .stderr_eq(file!["stderr.term.svg"]);
 

tests/testsuite/cargo_add/script_bare/out/cargo-test-fixture.rs

@@ -1 +1,6 @@
+---
+[dependencies]
+my-package = "99999.0.0"
+---
+
 fn main() {}

tests/testsuite/cargo_add/script_bare/stderr.term.svg

@@ -31,7 +31,7 @@ fn case() {
         .arg_line("--manifest-path cargo-test-fixture.rs my-package")
         .current_dir(cwd)
         .assert()
-        .failure()
+        .success()
         .stdout_eq(str![""])
         .stderr_eq(file!["stderr.term.svg"]);
 

tests/testsuite/cargo_add/script_frontmatter/mod.rs

@@ -1,6 +1,9 @@
 ---
 [package]
 edition = "2015"
+
+[dependencies]
+my-package = "99999.0.0"
 ---
 
 fn main() {

tests/testsuite/cargo_add/script_frontmatter/out/cargo-test-fixture.rs

@@ -31,7 +31,7 @@ fn case() {
         .arg_line("--manifest-path cargo-test-fixture.rs my-package")
         .current_dir(cwd)
         .assert()
-        .failure()
+        .success()
         .stdout_eq(str![""])
         .stderr_eq(file!["stderr.term.svg"]);
 

tests/testsuite/cargo_add/script_frontmatter/stderr.term.svg

@@ -1,3 +1,7 @@
 #!/usr/bin/env cargo
+---
+[dependencies]
+my-package = "99999.0.0"
+---
 
 fn main() {}