Typing overhaul

immutables/__init__.py

@@ -1,9 +1,14 @@
 # flake8: noqa
 
+from typing import TYPE_CHECKING
+
 try:
     from ._map import Map
 except ImportError:
-    from .map import Map
+    if TYPE_CHECKING:
+        from ._map import Map
+    else:
+        from .map import Map
 else:
     import collections.abc as _abc
     _abc.Mapping.register(Map)

immutables/_map.pyi

@@ -11,9 +11,11 @@ from typing import Tuple
 from typing import Type
 from typing import TypeVar
 from typing import Union
+from typing import Optional
 
 
 K = TypeVar('K', bound=Hashable)
+C = TypeVar('C', bound=Hashable)
 V = TypeVar('V', bound=Any)
 D = TypeVar('D', bound=Any)
 
@@ -41,26 +43,51 @@ class MapItems(Generic[K, V]):
 
 class Map(Mapping[K, V]):
     @overload
-    def __init__(self, **kw: V) -> None: ...
+    def __init__(self: Map[Any, Any]) -> None: ...
+
+    @overload
+    def __init__(self: Map[str, V], **kw: V) -> None: ...
+
     @overload
     def __init__(
-        self, col: Union[Mapping[K, V], Iterable[Tuple[K, V]]], **kw: V
+        self, col: Union[Mapping[K, V], Iterable[Tuple[K, V]]]
     ) -> None: ...
+    @overload
+    def __init__(
+        self: Map[Union[C, str], V],
+        col: Union[Mapping[C, V], Iterable[Tuple[C, V]]],
+        **kw: V
+    ) -> None: ...
+
     def __reduce__(self) -> Tuple[Type[Map], Tuple[dict]]: ...
     def __len__(self) -> int: ...
     def __eq__(self, other: Any) -> bool: ...
+
+    @overload
+    def update(self, **kw: D) -> Map[Union[K, str], Union[V, D]]: ...
+    # The overloads are overlapping, but not really.
+    # A call with no kwargs doesn't add string keys, so
+    # Map({1: "a"}).update({2: "b"}) produces Map[int, str], but
+    # Map({1: "a"}).update(two="b") produces Map[int|str, str]
     @overload
-    def update(self, **kw: V) -> Map[str, V]: ...
+    def update( # type: ignore
+        self,
+        col: Union[Mapping[C, D], Iterable[Tuple[C, D]]]
+    ) -> Map[Union[K, C], Union[V, D]]: ...
     @overload
     def update(
-        self, col: Union[Mapping[K, V], Iterable[Tuple[K, V]]], **kw: V
-    ) -> Map[K, V]: ...
+        self,
+        col: Union[Mapping[C, D], Iterable[Tuple[C, D]]],
+        **kw: D
+    ) -> Map[Union[K, C, str], Union[V, D]]: ...
+
     def mutate(self) -> MapMutation[K, V]: ...
-    def set(self, key: K, val: V) -> Map[K, V]: ...
+    def set(self, key: C, val: D) -> Map[Union[K, C], Union[V, D]]: ...
     def delete(self, key: K) -> Map[K, V]: ...
+    def get(self, key: K) -> Optional[V]: ...
     def get(self, key: K, default: D = ...) -> Union[V, D]: ...
     def __getitem__(self, key: K) -> V: ...
-    def __contains__(self, key: object) -> bool: ...
+    def __contains__(self, key: Any) -> bool: ...
     def __iter__(self) -> Iterator[K]: ...
     def keys(self) -> MapKeys[K]: ...
     def values(self) -> MapValues[V]: ...
@@ -86,10 +113,16 @@ class MapMutation(MutableMapping[K, V]):
     def __getitem__(self, key: K) -> V: ...
     def __contains__(self, key: Any) -> bool: ...
     @overload
-    def update(self, **kw: V) -> None: ...
+    def update(self: MapMutation[Union[C, str], V], **kw: V) -> None: ...
+    @overload
+    def update(
+        self, col: Union[Mapping[K, V], Iterable[Tuple[K, V]]]
+    ) -> None: ...
     @overload
     def update(
-        self, col: Union[Mapping[K, V], Iterable[Tuple[K, V]]], **kw: V
+        self: MapMutation[Union[C, str], V],
+        col: Union[Mapping[K, V], Iterable[Tuple[K, V]]],
+        **kw: V
     ) -> None: ...
     def finish(self) -> Map[K, V]: ...
     def __len__(self) -> int: ...