feat: Add TableLike::entry

toml-rs · Apr 26, 2022 · 6db12fd · 6db12fd
1 parent 4156603
commit 6db12fd
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 3 deletions.
diff --git a/src/inline_table.rs b/src/inline_table.rs
@@ -398,6 +398,29 @@ impl TableLike for InlineTable {
     fn clear(&mut self) {
         self.clear();
     }
+    fn entry<'a>(&'a mut self, key: &str) -> crate::Entry<'a> {
+        // Accept a `&str` rather than an owned type to keep `InternalString`, well, internal
+        match self.items.entry(key.into()) {
+            indexmap::map::Entry::Occupied(entry) => {
+                crate::Entry::Occupied(crate::OccupiedEntry { entry })
+            }
+            indexmap::map::Entry::Vacant(entry) => {
+                crate::Entry::Vacant(crate::VacantEntry { entry, key: None })
+            }
+        }
+    }
+    fn entry_format<'a>(&'a mut self, key: &Key) -> crate::Entry<'a> {
+        // Accept a `&Key` to be consistent with `entry`
+        match self.items.entry(key.get().into()) {
+            indexmap::map::Entry::Occupied(entry) => {
+                crate::Entry::Occupied(crate::OccupiedEntry { entry })
+            }
+            indexmap::map::Entry::Vacant(entry) => crate::Entry::Vacant(crate::VacantEntry {
+                entry,
+                key: Some(key.to_owned()),
+            }),
+        }
+    }
     fn get<'s>(&'s self, key: &str) -> Option<&'s Item> {
         self.items.get(key).map(|kv| &kv.value)
     }

diff --git a/src/table.rs b/src/table.rs
@@ -471,6 +471,10 @@ pub trait TableLike: crate::private::Sealed {
     }
     /// Clears the table, removing all key-value pairs. Keeps the allocated memory for reuse.
     fn clear(&mut self);
+    /// Gets the given key's corresponding entry in the Table for in-place manipulation.
+    fn entry<'a>(&'a mut self, key: &str) -> Entry<'a>;
+    /// Gets the given key's corresponding entry in the Table for in-place manipulation.
+    fn entry_format<'a>(&'a mut self, key: &Key) -> Entry<'a>;
     /// Returns an optional reference to an item given the key.
     fn get<'s>(&'s self, key: &str) -> Option<&'s Item>;
     /// Returns an optional mutable reference to an item given the key.
@@ -514,6 +518,12 @@ impl TableLike for Table {
     fn clear(&mut self) {
         self.clear();
     }
+    fn entry<'a>(&'a mut self, key: &str) -> Entry<'a> {
+        self.entry(key)
+    }
+    fn entry_format<'a>(&'a mut self, key: &Key) -> Entry<'a> {
+        self.entry_format(key)
+    }
     fn get<'s>(&'s self, key: &str) -> Option<&'s Item> {
         self.get(key)
     }
@@ -602,7 +612,7 @@ impl<'a> Entry<'a> {
 
 /// A view into a single occupied location in a `IndexMap`.
 pub struct OccupiedEntry<'a> {
-    entry: indexmap::map::OccupiedEntry<'a, InternalString, TableKeyValue>,
+    pub(crate) entry: indexmap::map::OccupiedEntry<'a, InternalString, TableKeyValue>,
 }
 
 impl<'a> OccupiedEntry<'a> {
@@ -656,8 +666,8 @@ impl<'a> OccupiedEntry<'a> {
 
 /// A view into a single empty location in a `IndexMap`.
 pub struct VacantEntry<'a> {
-    entry: indexmap::map::VacantEntry<'a, InternalString, TableKeyValue>,
-    key: Option<Key>,
+    pub(crate) entry: indexmap::map::VacantEntry<'a, InternalString, TableKeyValue>,
+    pub(crate) key: Option<Key>,
 }
 
 impl<'a> VacantEntry<'a> {