Repository: MagicStack/immutables

Issue:
mypy integration
Collections should implement [mypy Mapping generic](https://mypy.readthedocs.io/en/stable/generics.html#defining-sub-classes-of-generic-classes) to be properly typed

Discussion:
1st1: > Oh right, so should we make the pure Python version also implement __class_get_item__ instead?

Yes, that's the idea.

> Because I believe Mapping.register(Map) is not sufficient for making Map[foo, bar] return Map?

True.  We only need the `Mapping.register()` part to make `isinstance(immumap, Mapping)` work.
antonagestam: @1st1 Gotcha, thanks for taking the time to explain! I'll try to update the PR soon :) 
antonagestam: I added fixes that addresses your comments as well as a test case. Perhaps the documentation should be updated with an example as well? 
1st1: Looks great! Only a few nits left.
antonagestam: I have a few questions on the mutation API.

Should Map.mutate() always return a MapMutation with it's own signature? E.g. `Map[A, B]().mutate() => MapMutation[A, B]`? I think that would make sense as then MapMutation.finish() would return the same type Map.

As for both MapMutation and Map should I add type hints for all methods or just the ones that are part of the public API? There's a few methods that always raise an exception, like MapMutation.__iter__(), should these have the return type `typing.NoReturn` or just not be hinted?

I'm also curious as to why MapMutation.pop() accepts a variadic *args and then raises TypeError if it's longer than 1. Why not use a non-variadic single `default` argument? :)
1st1: > Should Map.mutate() always return a MapMutation with it's own signature? E.g. Map[A, B]().mutate() => MapMutation[A, B]?

Yes.

> As for both MapMutation and Map should I add type hints for all methods or just the ones that are part of the public API? 

Ideally, yes.

> There's a few methods that always raise an exception, like MapMutation.__iter__(), should these have the return type typing.NoReturn or just not be hinted?

I guess using `typing.NoReturn` makes sense.

> I'm also curious as to why MapMutation.pop() accepts a variadic *args and then raises TypeError if it's longer than 1. Why not use a non-variadic single default argument? :)

This is to emulate a "positional-only argument", IIRC. 
antonagestam: I added more type hints, so the mutation API should now be covered as well as the keys, values and items methods. I added a `BitmapNode` as an empty class because it is exposed as the type of the `m` arguments in `__init__` methods. However, it's implementation was not to decipher and so I left it untyped for now. I don't think I'll be able to figure out the correct type hints for it.
1st1: @ambv
antonagestam: Ping me if you need me to rebase or make other changes :)

Patch:
diff --git a/immutables/_map.c b/immutables/_map.c
index f9c8772..7e63562 100644
--- a/immutables/_map.c
+++ b/immutables/_map.c
@@ -3363,6 +3363,12 @@ map_reduce(MapObject *self)
     return tup;
 }
 
+static PyObject *
+map_py_class_getitem(PyObject *type, PyObject *item)
+{
+    Py_INCREF(type);
+    return type;
+}
 
 static PyMethodDef Map_methods[] = {
     {"set", (PyCFunction)map_py_set, METH_VARARGS, NULL},
@@ -3375,6 +3381,12 @@ static PyMethodDef Map_methods[] = {
     {"update", (PyCFunction)map_py_update, METH_VARARGS | METH_KEYWORDS, NULL},
     {"__reduce__", (PyCFunction)map_reduce, METH_NOARGS, NULL},
     {"__dump__", (PyCFunction)map_py_dump, METH_NOARGS, NULL},
+    {
+        "__class_getitem__",
+        (PyCFunction)map_py_class_getitem,
+        METH_O|METH_CLASS,
+        NULL
+    },
     {NULL, NULL}
 };
 
diff --git a/immutables/_map.pyi b/immutables/_map.pyi
new file mode 100644
index 0000000..863d911
--- /dev/null
+++ b/immutables/_map.pyi
@@ -0,0 +1,87 @@
+from typing import Any
+from typing import Generic
+from typing import Hashable
+from typing import Iterable
+from typing import Iterator
+from typing import Literal
+from typing import Mapping
+from typing import MutableMapping
+from typing import NoReturn
+from typing import Tuple
+from typing import Type
+from typing import TypeVar
+from typing import Union
+
+
+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]]: ...
+
+
+class Map(Mapping[K, V]):
+    def __init__(
+        self, col: Union[Mapping[K, V], Iterable[Tuple[K, V]]] = ..., **kw: V
+    ): ...
+    def __reduce__(self) -> NoReturn: ...
+    def __len__(self) -> int: ...
+    def __eq__(self, other: Any) -> bool: ...
+    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) -> Literal[False]: ...
+    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: ...
+    def update(
+        self, col: Union[Mapping[K, V], Iterable[Tuple[K, V]]] = ..., **kw: V
+    ): ...
+    def finish(self) -> Map[K, V]: ...
+    def __len__(self) -> int: ...
+    def __eq__(self, other: Any) -> bool: ...
diff --git a/immutables/map.py b/immutables/map.py
index 4c5cd49..3ea4656 100644
--- a/immutables/map.py
+++ b/immutables/map.py
@@ -630,6 +630,9 @@ class Map:
         self.__root.dump(buf, 0)
         return '\n'.join(buf)
 
+    def __class_getitem__(cls, item):
+        return cls
+
 
 class MapMutation:
 
diff --git a/immutables/py.typed b/immutables/py.typed
new file mode 100644
index 0000000..b648ac9
--- /dev/null
+++ b/immutables/py.typed
@@ -0,0 +1,1 @@
+partial


Test patch:
diff --git a/tests/test_map.py b/tests/test_map.py
index a99b856..8d629f8 100644
--- a/tests/test_map.py
+++ b/tests/test_map.py
@@ -2,6 +2,7 @@ import collections.abc
 import gc
 import pickle
 import random
+import sys
 import unittest
 import weakref
 
@@ -1296,6 +1297,10 @@ class BaseMapTest:
         with self.assertRaisesRegex(TypeError, "can't pickle"):
             pickle.dumps(h.mutate())
 
+    @unittest.skipIf(sys.version_info < (3, 7, 0), "__class_getitem__ is not available")
+    def test_map_is_subscriptable(self):
+        self.assertIs(self.Map[int, str], self.Map)
+
 
 class PyMapTest(BaseMapTest, unittest.TestCase):