diff --git a/docs/changelog.md b/docs/changelog.md index ed4b6d8330..529d929a3e 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -18,6 +18,10 @@ what we publish. ### ⭐️ New +- `Dict` now support `popitem`, which removes and returns the last item in the `Dict`. +([PR #2701](https://github.com/modularml/mojo/pull/2701) +by [@jayzhan211](https://github.com/jayzhan211)) + - Add a `sort` function for list of `ComparableCollectionElement`s. [PR #2609](https://github.com/modularml/mojo/pull/2609) by [@mzaks](https://github.com/mzaks) diff --git a/stdlib/src/collections/dict.mojo b/stdlib/src/collections/dict.mojo index 9d348ecc04..acd0bfdaf4 100644 --- a/stdlib/src/collections/dict.mojo +++ b/stdlib/src/collections/dict.mojo @@ -799,6 +799,34 @@ struct Dict[K: KeyElement, V: CollectionElement]( return default.value()[] raise "KeyError" + fn popitem(inout self) raises -> DictEntry[K, V]: + """Remove and return a (key, value) pair from the dictionary. Pairs are returned in LIFO order. + popitem() is useful to destructively iterate over a dictionary, as often used in set algorithms. + If the dictionary is empty, calling popitem() raises a KeyError. + + Args: None + + Returns: + Last dictionary item + + Raises: + "KeyError" if the dictionary is empty. + """ + + var key = Optional[K](None) + var val = Optional[V](None) + + for item in reversed(self.items()): + key = Optional(item[].key) + val = Optional(item[].value) + break + + if key: + _ = self.pop(key.value()[]) + return DictEntry[K, V](key.value()[], val.value()[]) + + raise "KeyError: popitem(): dictionary is empty" + fn keys( self: Reference[Self, _, _] ) -> _DictKeyIter[K, V, self.is_mutable, self.lifetime]: diff --git a/stdlib/test/collections/test_dict.mojo b/stdlib/test/collections/test_dict.mojo index b51f0cb13b..0788e4add1 100644 --- a/stdlib/test/collections/test_dict.mojo +++ b/stdlib/test/collections/test_dict.mojo @@ -452,6 +452,7 @@ def test_dict(): test["test_dict_update_empty_new", test_dict_update_empty_new]() test["test_mojo_issue_1729", test_mojo_issue_1729]() test["test dict or", test_dict_or]() + test["test dict popteim", test_dict_popitem]() def test_taking_owned_kwargs_dict(owned kwargs: OwnedKwargsDict[Int]): @@ -512,6 +513,21 @@ def test_find_get(): assert_equal(some_dict.get("not_key", 0), 0) +def test_dict_popitem(): + var dict = Dict[String, Int]() + dict["a"] = 1 + dict["b"] = 2 + + var item = dict.popitem() + assert_equal(item.key, "b") + assert_equal(item.value, 2) + item = dict.popitem() + assert_equal(item.key, "a") + assert_equal(item.value, 1) + with assert_raises(contains="KeyError"): + _ = dict.popitem() + + fn test_clear() raises: var some_dict = Dict[String, Int]() some_dict["key"] = 1