Refactor typings

* Improve typing of `__init__()`
* Update typing of `Map`-producing functions to produce the correct type
* Update typing of other methods to more closely align with `Mapping`
* Add protocol classes for unexposed data structures
* Export protocol classes for ease of use in typed code
* Update stub file to pass in mypy strict mode

8 files changed, 170 insertions, 80 deletions

diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 971e850..7558539 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -51,4 +51,5 @@ jobs:
       run: |
         pip install -e .[test]
         flake8 immutables/ tests/
+        mypy immutables/
         python -m unittest -v tests.suite
diff --git a/.gitignore b/.gitignore
index 4975f93..07d9825 100644
--- a/.gitignore
+++ b/.gitignore
@@ -20,3 +20,4 @@ __pycache__/
 /.pytest_cache
 /.coverage
 /.mypy_cache
+/.venv*
diff --git a/MANIFEST.in b/MANIFEST.in
index d3f372b..9d9b332 100644
--- a/MANIFEST.in
+++ b/MANIFEST.in
@@ -1,4 +1,4 @@
 recursive-include tests *.py
 recursive-include immutables *.py *.c *.h *.pyi
 include LICENSE* NOTICE README.rst bench.png
-include immutables/py.typed
+include mypy.ini immutables/py.typed
diff --git a/immutables/__init__.py b/immutables/__init__.py
index eaf0230..fc7764b 100644
--- a/immutables/__init__.py
+++ b/immutables/__init__.py
@@ -1,12 +1,27 @@
 # flake8: noqa
 
-try:
+import sys
+
+if sys.version_info >= (3, 5, 2):
+    from typing import TYPE_CHECKING
+else:
+    from typing_extensions import TYPE_CHECKING
+
+if TYPE_CHECKING:
     from ._map import Map
-except ImportError:
-    from .map import Map
 else:
-    import collections.abc as _abc
-    _abc.Mapping.register(Map)
+    try:
+        from ._map import Map
+    except ImportError:
+        from .map import Map
+    else:
+        import collections.abc as _abc
+        _abc.Mapping.register(Map)
+
+from ._protocols import MapKeys as MapKeys
+from ._protocols import MapValues as MapValues
+from ._protocols import MapItems as MapItems
+from ._protocols import MapMutation as MapMutation
 
 from ._version import __version__
 
diff --git a/immutables/_map.pyi b/immutables/_map.pyi
index b8a6de3..4590cd4 100644
--- a/immutables/_map.pyi
+++ b/immutables/_map.pyi
@@ -1,96 +1,73 @@
 from typing import Any
+from typing import Dict
 from typing import Generic
-from typing import Hashable
 from typing import Iterable
 from typing import Iterator
 from typing import Mapping
-from typing import MutableMapping
-from typing import NoReturn
-from typing import overload
+from typing import Optional
 from typing import Tuple
 from typing import Type
-from typing import TypeVar
 from typing import Union
+from typing import overload
 
-
-K = TypeVar('K', bound=Hashable)
-V = TypeVar('V', bound=Any)
-D = TypeVar('D', bound=Any)
-
-
-class BitmapNode: ...
-
-
-class MapKeys(Generic[K]):
-    def __init__(self, c: int, m: BitmapNode) -> None: ...
-    def __len__(self) -> int: ...
-    def __iter__(self) -> Iterator[K]: ...
-
-
-class MapValues(Generic[V]):
-    def __init__(self, c: int, m: BitmapNode) -> None: ...
-    def __len__(self) -> int: ...
-    def __iter__(self) -> Iterator[V]: ...
-
-
-class MapItems(Generic[K, V]):
-    def __init__(self, c: int, m: BitmapNode) -> None: ...
-    def __len__(self) -> int: ...
-    def __iter__(self) -> Iterator[Tuple[K, V]]: ...
+from ._protocols import IterableItems
+from ._protocols import MapItems
+from ._protocols import MapKeys
+from ._protocols import MapMutation
+from ._protocols import MapValues
+from ._protocols import HT
+from ._protocols import KT
+from ._protocols import T
+from ._protocols import VT_co
 
 
-class Map(Mapping[K, V]):
+class Map(Mapping[KT, VT_co]):
     @overload
-    def __init__(self, **kw: V) -> None: ...
+    def __init__(self) -> None: ...
+    @overload
+    def __init__(self: Map[str, VT_co], **kw: VT_co) -> None: ...
+    @overload
+    def __init__(
+        self, __col: Union[IterableItems[KT, VT_co], Iterable[Tuple[KT, VT_co]]]
+    ) -> None: ...
     @overload
     def __init__(
-        self, col: Union[Mapping[K, V], Iterable[Tuple[K, V]]], **kw: V
+        self: Map[Union[KT, str], VT_co],
+        __col: Union[IterableItems[KT, VT_co], Iterable[Tuple[KT, VT_co]]],
+        **kw: VT_co
     ) -> None: ...
-    def __reduce__(self) -> Tuple[Type[Map], Tuple[dict]]: ...
+    def __reduce__(self) -> Tuple[Type[Map[KT, VT_co]], Tuple[Dict[KT, VT_co]]]: ...
     def __len__(self) -> int: ...
     def __eq__(self, other: Any) -> bool: ...
     @overload
-    def update(self, **kw: V) -> Map[str, V]: ...
-    @overload
     def update(
-        self, col: Union[Mapping[K, V], Iterable[Tuple[K, V]]], **kw: V
-    ) -> Map[K, V]: ...
-    def mutate(self) -> MapMutation[K, V]: ...
-    def set(self, key: K, val: V) -> Map[K, V]: ...
-    def delete(self, key: K) -> Map[K, V]: ...
-    def get(self, key: K, default: D = ...) -> Union[V, D]: ...
-    def __getitem__(self, key: K) -> V: ...
-    def __contains__(self, key: object) -> bool: ...
-    def __iter__(self) -> Iterator[K]: ...
-    def keys(self) -> MapKeys[K]: ...
-    def values(self) -> MapValues[V]: ...
-    def items(self) -> MapItems[K, V]: ...
-    def __hash__(self) -> int: ...
-    def __dump__(self) -> str: ...
-    def __class_getitem__(cls, item: Any) -> Type[Map]: ...
-
-
-S = TypeVar('S', bound='MapMutation')
-
-
-class MapMutation(MutableMapping[K, V]):
-    def __init__(self, count: int, root: BitmapNode) -> None: ...
-    def set(self, key: K, val: V) -> None: ...
-    def __enter__(self: S) -> S: ...
-    def __exit__(self, *exc: Any): ...
-    def __iter__(self) -> NoReturn: ...
-    def __delitem__(self, key: K) -> None: ...
-    def __setitem__(self, key: K, val: V) -> None: ...
-    def pop(self, __key: K, __default: D = ...) -> Union[V, D]: ...
-    def get(self, key: K, default: D = ...) -> Union[V, D]: ...
-    def __getitem__(self, key: K) -> V: ...
-    def __contains__(self, key: Any) -> bool: ...
+        self,
+        __col: Union[IterableItems[KT, VT_co], Iterable[Tuple[KT, VT_co]]]
+    ) -> Map[KT, VT_co]: ...
     @overload
-    def update(self, **kw: V) -> None: ...
+    def update(
+        self: Map[Union[HT, str], Any],
+        __col: Union[IterableItems[KT, VT_co], Iterable[Tuple[KT, VT_co]]],
+        **kw: VT_co  # type: ignore[misc]
+    ) -> Map[KT, VT_co]: ...
     @overload
     def update(
-        self, col: Union[Mapping[K, V], Iterable[Tuple[K, V]]], **kw: V
-    ) -> None: ...
-    def finish(self) -> Map[K, V]: ...
-    def __len__(self) -> int: ...
-    def __eq__(self, other: Any) -> bool: ...
+        self: Map[Union[HT, str], Any],
+        **kw: VT_co  # type: ignore[misc]
+    ) -> Map[KT, VT_co]: ...
+    def mutate(self) -> MapMutation[KT, VT_co]: ...
+    def set(self, key: KT, val: VT_co) -> Map[KT, VT_co]: ...  # type: ignore[misc]
+    def delete(self, key: KT) -> Map[KT, VT_co]: ...
+    @overload
+    def get(self, key: KT) -> Optional[VT_co]: ...
+    @overload
+    def get(self, key: KT, default: Union[VT_co, T]) -> Union[VT_co, T]: ...
+    def __getitem__(self, key: KT) -> VT_co: ...
+    def __contains__(self, key: Any) -> bool: ...
+    def __iter__(self) -> Iterator[KT]: ...
+    def keys(self) -> MapKeys[KT]: ...  # type: ignore[override]
+    def values(self) -> MapValues[VT_co]: ...  # type: ignore[override]
+    def items(self) -> MapItems[KT, VT_co]: ...  # type: ignore[override]
+    def __hash__(self) -> int: ...
+    def __dump__(self) -> str: ...
+    def __class_getitem__(cls, item: Any) -> Type[Map[Any, Any]]: ...
diff --git a/immutables/_protocols.py b/immutables/_protocols.py
new file mode 100644
index 0000000..de87d23
--- /dev/null
+++ b/immutables/_protocols.py
@@ -0,0 +1,85 @@
+import sys
+from typing import Any
+from typing import Hashable
+from typing import Iterable
+from typing import Iterator
+from typing import NoReturn
+from typing import Optional
+from typing import Tuple
+from typing import TypeVar
+from typing import Union
+from typing import overload
+
+if sys.version_info >= (3, 8):
+    from typing import Protocol
+    from typing import TYPE_CHECKING
+else:
+    from typing_extensions import Protocol
+    from typing_extensions import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from ._map import Map
+
+HT = TypeVar('HT', bound=Hashable)
+KT = TypeVar('KT', bound=Hashable)
+KT_co = TypeVar('KT_co', covariant=True)
+MM = TypeVar('MM', bound='MapMutation[Any, Any]')
+T = TypeVar('T')
+VT = TypeVar('VT')
+VT_co = TypeVar('VT_co', covariant=True)
+
+
+class MapKeys(Protocol[KT_co]):
+    def __len__(self) -> int: ...
+    def __iter__(self) -> Iterator[KT_co]: ...
+
+
+class MapValues(Protocol[VT_co]):
+    def __len__(self) -> int: ...
+    def __iter__(self) -> Iterator[VT_co]: ...
+
+
+class MapItems(Protocol[KT_co, VT_co]):
+    def __len__(self) -> int: ...
+    def __iter__(self) -> Iterator[Tuple[KT_co, VT_co]]: ...
+
+
+class IterableItems(Protocol[KT_co, VT_co]):
+    def items(self) -> Iterable[Tuple[KT_co, VT_co]]: ...
+
+
+class MapMutation(Protocol[KT, VT]):
+    def set(self, key: KT, val: VT) -> None: ...
+    def __enter__(self: MM) -> MM: ...
+    def __exit__(self, *exc: Any) -> bool: ...
+    def __iter__(self) -> NoReturn: ...
+    def __delitem__(self, key: KT) -> None: ...
+    def __setitem__(self, key: KT, val: VT) -> None: ...
+    @overload
+    def pop(self, __key: KT) -> VT: ...
+    @overload
+    def pop(self, __key: KT, __default: T) -> Union[VT, T]: ...
+    @overload
+    def get(self, key: KT) -> Optional[VT]: ...
+    @overload
+    def get(self, key: KT, default: Union[VT, T]) -> Union[VT, T]: ...
+    def __getitem__(self, key: KT) -> VT: ...
+    def __contains__(self, key: object) -> bool: ...
+
+    @overload
+    def update(
+        self,
+        __col: Union[IterableItems[KT, VT], Iterable[Tuple[KT, VT]]]
+    ) -> None: ...
+
+    @overload
+    def update(
+        self: 'MapMutation[Union[HT, str], Any]',
+        __col: Union[IterableItems[KT, VT], Iterable[Tuple[KT, VT]]],
+        **kw: VT
+    ) -> None: ...
+    @overload
+    def update(self: 'MapMutation[Union[HT, str], Any]', **kw: VT) -> None: ...
+    def finish(self) -> 'Map[KT, VT]': ...
+    def __len__(self) -> int: ...
+    def __eq__(self, other: Any) -> bool: ...
diff --git a/mypy.ini b/mypy.ini
new file mode 100644
index 0000000..6b604b8
--- /dev/null
+++ b/mypy.ini
@@ -0,0 +1,9 @@
+[mypy]
+incremental = True
+strict = True
+
+[mypy-immutables.map]
+ignore_errors = True
+
+[mypy-immutables._testutils]
+ignore_errors = True
diff --git a/setup.py b/setup.py
index ba8e1da..530013c 100644
--- a/setup.py
+++ b/setup.py
@@ -10,6 +10,7 @@ TEST_DEPENDENCIES = [
     # (example breakage: https://gitlab.com/pycqa/flake8/issues/427)
     'flake8~=3.8.4',
     'pycodestyle~=2.6.0',
+    'mypy>=0.800',
 ]
 
 EXTRA_DEPENDENCIES = {
@@ -86,5 +87,6 @@ setuptools.setup(
     provides=['immutables'],
     include_package_data=True,
     ext_modules=ext_modules,
+    install_requires=['typing-extensions>=3.7.4.3;python_version<"3.8"'],
     extras_require=EXTRA_DEPENDENCIES,
 )