Merge pull request #788 from epage/encode-key

fix(edit): Preserve new key's formatting when inserting
toml-rs · Sep 24, 2024 · 5ae07ee · 5ae07ee
2 parents e71fc23 + 9bca304
commit 5ae07ee
Show file tree

Hide file tree

Showing 3 changed files with 95 additions and 11 deletions.
diff --git a/crates/toml_edit/src/inline_table.rs b/crates/toml_edit/src/inline_table.rs
@@ -381,20 +381,37 @@ impl InlineTable {
 
     /// Inserts a key-value pair into the map.
     pub fn insert(&mut self, key: impl Into<InternalString>, value: Value) -> Option<Value> {
-        let key = Key::new(key.into());
+        use indexmap::map::MutableEntryKey;
+        let key = Key::new(key);
         let value = Item::Value(value);
-        self.items
-            .insert(key, value)
-            .and_then(|old| old.into_value().ok())
+        match self.items.entry(key.clone()) {
+            indexmap::map::Entry::Occupied(mut entry) => {
+                entry.key_mut().fmt();
+                let old = std::mem::replace(entry.get_mut(), value);
+                old.into_value().ok()
+            }
+            indexmap::map::Entry::Vacant(entry) => {
+                entry.insert(value);
+                None
+            }
+        }
     }
 
     /// Inserts a key-value pair into the map.
     pub fn insert_formatted(&mut self, key: &Key, value: Value) -> Option<Value> {
-        let key = key.to_owned();
+        use indexmap::map::MutableEntryKey;
         let value = Item::Value(value);
-        self.items
-            .insert(key, value)
-            .and_then(|old| old.into_value().ok())
+        match self.items.entry(key.clone()) {
+            indexmap::map::Entry::Occupied(mut entry) => {
+                *entry.key_mut() = key.clone();
+                let old = std::mem::replace(entry.get_mut(), value);
+                old.into_value().ok()
+            }
+            indexmap::map::Entry::Vacant(entry) => {
+                entry.insert(value);
+                None
+            }
+        }
     }
 
     /// Removes an item given the key.

diff --git a/crates/toml_edit/src/table.rs b/crates/toml_edit/src/table.rs
@@ -393,14 +393,35 @@ impl Table {
 
     /// Inserts a key-value pair into the map.
     pub fn insert(&mut self, key: &str, item: Item) -> Option<Item> {
+        use indexmap::map::MutableEntryKey;
         let key = Key::new(key);
-        self.items.insert(key, item)
+        match self.items.entry(key.clone()) {
+            indexmap::map::Entry::Occupied(mut entry) => {
+                entry.key_mut().fmt();
+                let old = std::mem::replace(entry.get_mut(), item);
+                Some(old)
+            }
+            indexmap::map::Entry::Vacant(entry) => {
+                entry.insert(item);
+                None
+            }
+        }
     }
 
     /// Inserts a key-value pair into the map.
     pub fn insert_formatted(&mut self, key: &Key, item: Item) -> Option<Item> {
-        let key = key.to_owned();
-        self.items.insert(key, item)
+        use indexmap::map::MutableEntryKey;
+        match self.items.entry(key.clone()) {
+            indexmap::map::Entry::Occupied(mut entry) => {
+                *entry.key_mut() = key.clone();
+                let old = std::mem::replace(entry.get_mut(), item);
+                Some(old)
+            }
+            indexmap::map::Entry::Vacant(entry) => {
+                entry.insert(item);
+                None
+            }
+        }
     }
 
     /// Removes an item given the key.

diff --git a/crates/toml_edit/tests/testsuite/edit.rs b/crates/toml_edit/tests/testsuite/edit.rs
@@ -980,3 +980,49 @@ fn sorting_with_references() {
     let mut array = toml_edit::Array::from_iter(values);
     array.sort_by(|lhs, rhs| lhs.as_str().cmp(&rhs.as_str()));
 }
+
+#[test]
+fn table_str_key_whitespace() {
+    let mut document = "bookmark = 1010".parse::<DocumentMut>().unwrap();
+
+    let key: &str = "bookmark";
+
+    document.insert(key, array());
+    let table = document[key].as_array_of_tables_mut().unwrap();
+
+    let mut bookmark_table = Table::new();
+    bookmark_table["name"] = value("test.swf".to_owned());
+    table.push(bookmark_table);
+
+    assert_data_eq!(
+        document.to_string(),
+        str![[r#"
+[[bookmark]]
+name = "test.swf"
+
+"#]]
+    );
+}
+
+#[test]
+fn table_key_decor_whitespace() {
+    let mut document = "bookmark = 1010".parse::<DocumentMut>().unwrap();
+
+    let key = Key::parse("  bookmark   ").unwrap().remove(0);
+
+    document.insert_formatted(&key, array());
+    let table = document[&key].as_array_of_tables_mut().unwrap();
+
+    let mut bookmark_table = Table::new();
+    bookmark_table["name"] = value("test.swf".to_owned());
+    table.push(bookmark_table);
+
+    assert_data_eq!(
+        document.to_string(),
+        str![[r#"
+[[  bookmark   ]]
+name = "test.swf"
+
+"#]]
+    );
+}