Add typing information (#13)

5 files changed, 108 insertions, 0 deletions

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 @@
+partial
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):