Type annotation and test for HashTrieMap.update.

crate-py · Feb 13, 2024 · ce3bf37 · ce3bf37
1 parent 7ff2d43
commit ce3bf37
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 3 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "rpds-py"
-version = "0.17.1"
+version = "0.18.0"
 edition = "2021"
 
 [lib]

diff --git a/rpds.pyi b/rpds.pyi
@@ -11,6 +11,8 @@ from typing import (
 _T = TypeVar("_T")
 _KT_co = TypeVar("_KT_co", covariant=True)
 _VT_co = TypeVar("_VT_co", covariant=True)
+_KU_co = TypeVar("_KU_co", covariant=True)
+_VU_co = TypeVar("_VU_co", covariant=True)
 
 class HashTrieMap(Mapping[_KT_co, _VT_co]):
     def __init__(
@@ -31,7 +33,10 @@ class HashTrieMap(Mapping[_KT_co, _VT_co]):
         key: _KT_co,
         val: _VT_co,
     ) -> HashTrieMap[_KT_co, _VT_co]: ...
-    def update(self, *args: Mapping): ...
+    def update(
+        self,
+        *args: Mapping[_KU_co, _VU_co] | Iterable[tuple[_KU_co, _VU_co]],
+    ) -> HashTrieMap[_KT_co | _KU_co, _VT_co | _VU_co]: ...
     @classmethod
     def convert(
         cls,

diff --git a/tests/test_hash_trie_map.py b/tests/test_hash_trie_map.py
@@ -535,3 +535,13 @@ def test_fromkeys_explicit_value_not_copied():
     got[3].append(1)
 
     assert got == HashTrieMap((i, [1]) for i in keys)
+
+
+def test_update_with_iterable_of_kvs():
+    assert HashTrieMap({1: 2}).update(iter([(3, 4), ("5", 6)])) == HashTrieMap(
+        {
+            1: 2,
+            3: 4,
+            "5": 6,
+        },
+    )