From ee5bcc7a93b1ca993d3df9d500c4f0101bb15735 Mon Sep 17 00:00:00 2001
From: Mark Nevill <mark.nevill@gmail.com>
Date: Fri, 19 Jan 2024 09:18:06 +0800
Subject: [PATCH] set: Add option to write back to file

---
 src/main.rs  | 22 +++++++++++++++++-----
 test/test.rs | 24 +++++++++++++++++++-----
 2 files changed, 36 insertions(+), 10 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 0f85a13..7199c81 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -42,7 +42,10 @@ enum Args {
         opts: GetOpts,
     },
 
-    /// Edit the file to set some data (currently, just print modified version)
+    /// Edit the file to set some data.
+    ///
+    /// Specify `--write`/`-w` to write back to the input file, otherwise the
+    /// modified file is printed.
     Set {
         /// Path to the TOML file to read
         #[structopt(parse(from_os_str))]
@@ -53,6 +56,10 @@ enum Args {
 
         /// String value to place at the given spot (bool, array, etc. are TODO)
         value_str: String, // TODO more forms
+
+        /// Write back to file
+        #[structopt(short, long)]
+        write: bool,
     },
     //
     // TODO: append/add (name TBD)
@@ -95,7 +102,8 @@ fn main() {
             path,
             query,
             value_str,
-        } => set(&path, &query, &value_str),
+            write,
+        } => set(&path, &query, &value_str, write),
     };
     result.unwrap_or_else(|err| {
         match err.downcast::<SilentError>() {
@@ -183,7 +191,7 @@ fn print_toml_fragment(doc: &Document, tpath: &[TpathSegment]) {
     print!("{}", doc);
 }
 
-fn set(path: &PathBuf, query: &str, value_str: &str) -> Result<(), Error> {
+fn set(path: &PathBuf, query: &str, value_str: &str, write: bool) -> Result<(), Error> {
     let tpath = parse_query_cli(query)?.0;
     let mut doc = read_parse(path)?;
 
@@ -229,8 +237,12 @@ fn set(path: &PathBuf, query: &str, value_str: &str) -> Result<(), Error> {
     }
     *item = value(value_str);
 
-    // TODO actually write back
-    print!("{}", doc);
+    if write {
+        fs::write(path, doc.to_string())?;
+    } else {
+        print!("{}", doc);
+    }
+
     Ok(())
 }
 
diff --git a/test/test.rs b/test/test.rs
index 2d705a1..3d1db7a 100644
--- a/test/test.rs
+++ b/test/test.rs
@@ -104,12 +104,18 @@ tomltest_get_err_empty!(get_missing, ["nosuchkey"]);
 tomltest_get_err_empty!(get_missing_num, ["key[1]"]);
 
 macro_rules! tomltest_set {
-    ($name:ident, $args:expr, $expected:expr) => {
+    ($name:ident, $name_w:ident, $args:expr, $expected:expr) => {
         tomltest!($name, |mut t: TestCaseState| {
             t.write_file(INITIAL);
             t.cmd.args(["set", &t.filename()]).args($args);
             check_eq(&$expected, &t.expect_success());
         });
+        tomltest!($name_w, |mut t: TestCaseState| {
+            t.write_file(INITIAL);
+            t.cmd.args(["set", "-w", &t.filename()]).args($args);
+            check_eq("", &t.expect_success());
+            check_eq(&$expected, &t.read_file());
+        });
     };
 }
 
@@ -119,25 +125,26 @@ y = 1
 "#;
 
 #[rustfmt::skip]
-tomltest_set!(set_string_existing, ["x.y", "new"], r#"
+tomltest_set!(set_string_existing, set_string_existing_and_write, ["x.y", "new"], r#"
 [x]
 y = "new"
 "#);
 
 #[rustfmt::skip]
-tomltest_set!(set_string_existing_table, ["x.z", "123"], format!(
+tomltest_set!(set_string_existing_table, set_string_existing_table_and_write, ["x.z", "123"], 
+format!(
 r#"{INITIAL}z = "123"
 "#));
 
 #[rustfmt::skip]
-tomltest_set!(set_string_new_table, ["foo.bar", "baz"], format!(
+tomltest_set!(set_string_new_table, set_string_new_table_and_write, ["foo.bar", "baz"], format!(
 r#"{INITIAL}
 [foo]
 bar = "baz"
 "#));
 
 #[rustfmt::skip]
-tomltest_set!(set_string_toplevel, ["foo", "bar"], format!(
+tomltest_set!(set_string_toplevel, set_string_toplevel_and_write, ["foo", "bar"], format!(
 r#"foo = "bar"
 {INITIAL}"#));
 
@@ -181,6 +188,13 @@ impl TestCaseState {
         String::from_utf8(out.stderr).unwrap()
     }
 
+    pub fn read_file(&self) -> String {
+        let data = fs::read(&self.filename).expect("failed to read test fixture");
+        str::from_utf8(data.as_slice())
+            .expect("test fixture was not valid utf-8")
+            .to_owned()
+    }
+
     fn fail(&self, out: &Output, summary: &str) {
         panic!(
             "\n============\