Add PEP 585 GenericAlias support

immutables/_map.c

@@ -3373,12 +3373,14 @@ map_reduce(MapObject *self)
     return tup;
 }
 
+#if PY_VERSION_HEX < 0x030900A6
 static PyObject *
 map_py_class_getitem(PyObject *type, PyObject *item)
 {
     Py_INCREF(type);
     return type;
 }
+#endif
 
 static PyMethodDef Map_methods[] = {
     {"set", (PyCFunction)map_py_set, METH_VARARGS, NULL},
@@ -3393,9 +3395,13 @@ static PyMethodDef Map_methods[] = {
     {"__dump__", (PyCFunction)map_py_dump, METH_NOARGS, NULL},
     {
         "__class_getitem__",
+#if PY_VERSION_HEX < 0x030900A6
         (PyCFunction)map_py_class_getitem,
+#else
+        Py_GenericAlias,
+#endif
         METH_O|METH_CLASS,
-        NULL
+        "See PEP 585"
     },
     {NULL, NULL}
 };

immutables/_map.pyi

@@ -1,3 +1,4 @@
+import sys
 from typing import Any
 from typing import Dict
 from typing import Generic
@@ -10,6 +11,9 @@ from typing import Type
 from typing import Union
 from typing import overload
 
+if sys.version_info >= (3, 9):
+    from types import GenericAlias
+
 from ._protocols import IterableItems
 from ._protocols import MapItems
 from ._protocols import MapKeys
@@ -70,4 +74,7 @@ class Map(Mapping[KT, VT_co]):
     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]]: ...
+    if sys.version_info >= (3, 9):
+        def __class_getitem__(cls, item: Any) -> GenericAlias: ...
+    else:
+        def __class_getitem__(cls, item: Any) -> Type[Map[Any, Any]]: ...

immutables/map.py

@@ -2,6 +2,7 @@
 import itertools
 import reprlib
 import sys
+import types
 
 
 __all__ = ('Map',)
@@ -661,8 +662,11 @@ def __dump__(self):  # pragma: no cover
         self.__root.dump(buf, 0)
         return '\n'.join(buf)
 
-    def __class_getitem__(cls, item):
-        return cls
+    if sys.version_info >= (3, 9):
+        __class_getitem__ = classmethod(types.GenericAlias)
+    else:
+        def __class_getitem__(cls, item):
+            return cls
 
 
 class MapMutation:

tests/test_map.py

@@ -1384,11 +1384,13 @@ def test_map_pickle(self):
         with self.assertRaisesRegex(TypeError, "can('t|not) 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)
+        if sys.version_info >= (3, 9):
+            with_args = self.Map[int, str]
+            self.assertIs(with_args.__origin__, self.Map)
+            self.assertEqual(with_args.__args__, (int, str))
+        else:
+            self.assertIs(self.Map[int, str], self.Map)
 
     def test_kwarg_named_col(self):
         self.assertEqual(dict(self.Map(col=0)), {"col": 0})