[stdlib] Implement Dict.setdefault(key, default)

docs/changelog.md

@@ -102,6 +102,11 @@ what we publish.
 
 - The `Reference` type (and many iterators) now use "inferred" parameters to
   represent the mutability of their lifetime, simplifying the interface.
+
+- `Dict` now implements `setdefault`, to get a value from the dictionary by
+  key, or set it to a default if it doesn't exist
+  ([PR #2803](https://github.com/modularml/mojo/pull/2803)
+  by [@msaelices](https://github.com/msaelices))
 
 - Added new `ExplicitlyCopyable` trait, to mark types that can be copied
   explicitly, but which might not be implicitly copyable.

stdlib/src/collections/dict.mojo

@@ -870,6 +870,24 @@ struct Dict[K: KeyElement, V: CollectionElement](
         self._entries = Self._new_entries(Self._initial_reservation)
         self._index = _DictIndex(self._reserved())
 
+    fn setdefault(
+        ref [_]self: Self, key: K, owned default: V
+    ) raises -> Reference[V, __lifetime_of(self)]:
+        """Get a value from the dictionary by key, or set it to a default if it doesn't exist.
+
+        Args:
+            key: The key to search for in the dictionary.
+            default: The default value to set if the key is not present.
+
+        Returns:
+            The value associated with the key, or the default value if it wasn't present.
+        """
+        try:
+            return self.__get_ref(key)
+        except KeyError:
+            self[key] = default^
+            return self.__get_ref(key)
+
     @staticmethod
     @always_inline
     fn _new_entries(reserve_at_least: Int) -> List[Optional[DictEntry[K, V]]]:

stdlib/test/collections/test_dict.mojo

@@ -548,10 +548,33 @@ fn test_clear() raises:
     assert_equal(len(some_dict), 0)
 
 
+fn test_dict_setdefault() raises:
+    var some_dict = Dict[String, Int]()
+    some_dict["key1"] = 1
+    some_dict["key2"] = 2
+    assert_equal(some_dict.setdefault("key1", 0)[], 1)
+    assert_equal(some_dict.setdefault("key2", 0)[], 2)
+    assert_equal(some_dict.setdefault("not_key", 0)[], 0)
+    assert_equal(some_dict["not_key"], 0)
+
+    # Check that there is no copy of the default value, so it's performant
+    var other_dict = Dict[String, CopyCounter]()
+    var a = CopyCounter()
+    var a_def = CopyCounter()
+    var b_def = CopyCounter()
+    other_dict["a"] = a^
+    assert_equal(1, other_dict["a"].copy_count)
+    _ = other_dict.setdefault("a", a_def^)
+    _ = other_dict.setdefault("b", b_def^)
+    assert_equal(1, other_dict["a"].copy_count)
+    assert_equal(1, other_dict["b"].copy_count)
+
+
 def main():
     test_dict()
     test_dict_fromkeys()
     test_dict_fromkeys_optional()
+    test_dict_setdefault()
     test_dict_string_representation_string_int()
     test_dict_string_representation_int_int()
     test_owned_kwargs_dict()