Add AsyncContextManager generic class (#438)

Asynchronous context managers are defined by PEP 492, but there is no corresponding generic abstract class in typing. This PR adds it for Python 3.5+.

Author: ilevkivskyi

src/test_typing.py

@@ -1552,6 +1552,12 @@ def __anext__(self) -> T_a:
             return data
         else:
             raise StopAsyncIteration
+
+class ACM:
+    async def __aenter__(self) -> int:
+        return 42
+    async def __aexit__(self, etype, eval, tb):
+        return None
 """
 
 if ASYNCIO:
@@ -1562,12 +1568,13 @@ def __anext__(self) -> T_a:
 else:
     # fake names for the sake of static analysis
     asyncio = None
-    AwaitableWrapper = AsyncIteratorWrapper = object
+    AwaitableWrapper = AsyncIteratorWrapper = ACM = object
 
 PY36 = sys.version_info[:2] >= (3, 6)
 
 PY36_TESTS = """
 from test import ann_module, ann_module2, ann_module3
+from typing import AsyncContextManager
 
 class A:
     y: float
@@ -1604,6 +1611,16 @@ def __str__(self):
         return f'{self.x} -> {self.y}'
     def __add__(self, other):
         return 0
+
+async def g_with(am: AsyncContextManager[int]):
+    x: int
+    async with am as x:
+        return x
+
+try:
+    g_with(ACM()).send(None)
+except StopIteration as e:
+    assert e.args[0] == 42
 """
 
 if PY36:
@@ -2165,6 +2182,24 @@ def manager():
         self.assertIsInstance(cm, typing.ContextManager)
         self.assertNotIsInstance(42, typing.ContextManager)
 
+    @skipUnless(ASYNCIO, 'Python 3.5 required')
+    def test_async_contextmanager(self):
+        class NotACM:
+            pass
+        self.assertIsInstance(ACM(), typing.AsyncContextManager)
+        self.assertNotIsInstance(NotACM(), typing.AsyncContextManager)
+        @contextlib.contextmanager
+        def manager():
+            yield 42
+
+        cm = manager()
+        self.assertNotIsInstance(cm, typing.AsyncContextManager)
+        self.assertEqual(typing.AsyncContextManager[int].__args__, (int,))
+        with self.assertRaises(TypeError):
+            isinstance(42, typing.AsyncContextManager[int])
+        with self.assertRaises(TypeError):
+            typing.AsyncContextManager[int, str]
+
 
 class TypeTests(BaseTestCase):
 

src/typing.py

@@ -10,6 +10,8 @@
     import collections.abc as collections_abc
 except ImportError:
     import collections as collections_abc  # Fallback for PY3.2.
+if sys.version_info[:2] >= (3, 6):
+    import _collections_abc  # Needed for private function _check_methods # noqa
 try:
     from types import WrapperDescriptorType, MethodWrapperType, MethodDescriptorType
 except ImportError:
@@ -59,6 +61,7 @@
     # Coroutine,
     # Collection,
     # AsyncGenerator,
+    # AsyncContextManager
 
     # Structural checks, a.k.a. protocols.
     'Reversible',
@@ -1974,6 +1977,38 @@ def __subclasshook__(cls, C):
             return NotImplemented
 
 
+if hasattr(contextlib, 'AbstractAsyncContextManager'):
+    class AsyncContextManager(Generic[T_co],
+                              extra=contextlib.AbstractAsyncContextManager):
+        __slots__ = ()
+
+    __all__.append('AsyncContextManager')
+elif sys.version_info[:2] >= (3, 5):
+    exec("""
+class AsyncContextManager(Generic[T_co]):
+    __slots__ = ()
+
+    async def __aenter__(self):
+        return self
+
+    @abc.abstractmethod
+    async def __aexit__(self, exc_type, exc_value, traceback):
+        return None
+
+    @classmethod
+    def __subclasshook__(cls, C):
+        if cls is AsyncContextManager:
+            if sys.version_info[:2] >= (3, 6):
+                return _collections_abc._check_methods(C, "__aenter__", "__aexit__")
+            if (any("__aenter__" in B.__dict__ for B in C.__mro__) and
+                    any("__aexit__" in B.__dict__ for B in C.__mro__)):
+                return True
+        return NotImplemented
+
+__all__.append('AsyncContextManager')
+""")
+
+
 class Dict(dict, MutableMapping[KT, VT], extra=dict):
 
     __slots__ = ()