fix(edit): Preserve new key's formatting when inserting

Fixes #787
toml-rs · Sep 24, 2024 · 9bca304 · 9bca304
1 parent c3600e6
commit 9bca304
Show file tree

Hide file tree

Showing 3 changed files with 51 additions and 13 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
@@ -997,7 +997,7 @@ fn table_str_key_whitespace() {
     assert_data_eq!(
         document.to_string(),
         str![[r#"
-[[bookmark ]]
+[[bookmark]]
 name = "test.swf"
 
 "#]]
@@ -1020,7 +1020,7 @@ fn table_key_decor_whitespace() {
     assert_data_eq!(
         document.to_string(),
         str![[r#"
-[[bookmark ]]
+[[  bookmark   ]]
 name = "test.swf"
 
 "#]]