Fixes to ContextManager (#1249)

* add typing.ContextManager for 3.6+ only

This fixes the easier part of #655.

Would it make sense to add a generic typing.ContextManager that exists in any Python version?

* update comment

* fix argument types for ContextManager.__exit__

* add AsyncContextManager

* add @asynccontextmanager

* typing.ContextManager now always exists

* back out async-related changes

Will submit those in a separate PR later

* fix import order

* AbstractContextManager only exists in 3.6+

* AbstractContextManager -> ContextManager

stdlib/2/typing.pyi

@@ -1,7 +1,7 @@
 # Stubs for typing (Python 2.7)
 
 from abc import abstractmethod, ABCMeta
-from types import CodeType, FrameType
+from types import CodeType, FrameType, TracebackType
 
 # Definitions of special type checking related constructs.  Their definitions
 # are not used, so their value does not matter.
@@ -199,6 +199,12 @@ class ValuesView(MappingView, Iterable[_VT_co], Generic[_VT_co]):
     def __contains__(self, o: object) -> bool: ...
     def __iter__(self) -> Iterator[_VT_co]: ...
 
+class ContextManager(Generic[_T_co]):
+    def __enter__(self) -> _T_co: ...
+    def __exit__(self, exc_type: Optional[Type[BaseException]],
+                 exc_value: Optional[BaseException],
+                 traceback: Optional[TracebackType]) -> Optional[bool]: ...
+
 class Mapping(Iterable[_KT], Container[_KT], Sized, Generic[_KT, _VT_co]):
     # TODO: We wish the key type could also be covariant, but that doesn't work,
     # see discussion in https: //github.com/python/typing/pull/273.

stdlib/2and3/contextlib.pyi

@@ -2,24 +2,23 @@
 
 from typing import (
     Any, Callable, Generator, IO, Iterable, Iterator, Optional, Type,
-    Generic, TypeVar,
+    Generic, TypeVar
 )
 from types import TracebackType
 import sys
+# Aliased here for backwards compatibility; TODO eventually remove this
+from typing import ContextManager as ContextManager
+
+if sys.version_info >= (3, 6):
+    from typing import ContextManager as AbstractContextManager
 
 _T = TypeVar('_T')
+
 _ExitFunc = Callable[[Optional[Type[BaseException]],
-                      Optional[Exception],
+                      Optional[BaseException],
                       Optional[TracebackType]], bool]
 _CM_EF = TypeVar('_CM_EF', ContextManager, _ExitFunc)
 
-# TODO already in PEP, have to get added to mypy
-class ContextManager(Generic[_T]):
-    def __enter__(self) -> _T: ...
-    def __exit__(self, exc_type: Optional[Type[BaseException]],
-                 exc_val: Optional[Exception],
-                 exc_tb: Optional[TracebackType]) -> bool: ...
-
 if sys.version_info >= (3, 2):
     class GeneratorContextManager(Generic[_T], ContextManager[_T]):
         def __call__(self, func: Callable[..., _T]) -> Callable[..., _T]: ...

stdlib/3/typing.pyi

@@ -2,7 +2,7 @@
 
 import sys
 from abc import abstractmethod, ABCMeta
-from types import CodeType, FrameType
+from types import CodeType, FrameType, TracebackType
 
 # Definitions of special type checking related constructs.  Their definition
 # are not used, so their value does not matter.
@@ -127,7 +127,7 @@ class Generator(Iterator[_T_co], Generic[_T_co, _T_contra, _V_co]):
     gi_yieldfrom = ...  # type: Optional[Generator]
 
 # TODO: Several types should only be defined if sys.python_version >= (3, 5):
-# Awaitable, AsyncIterator, AsyncIterable, Coroutine, Collection, ContextManager.
+# Awaitable, AsyncIterator, AsyncIterable, Coroutine, Collection.
 # See https: //github.com/python/typeshed/issues/655 for why this is not easy.
 
 class Awaitable(Generic[_T_co]):
@@ -281,7 +281,11 @@ class ValuesView(MappingView, Iterable[_VT_co], Generic[_VT_co]):
     def __contains__(self, o: object) -> bool: ...
     def __iter__(self) -> Iterator[_VT_co]: ...
 
-# TODO: ContextManager (only if contextlib.AbstractContextManager exists)
+class ContextManager(Generic[_T_co]):
+    def __enter__(self) -> _T_co: ...
+    def __exit__(self, exc_type: Optional[Type[BaseException]],
+                 exc_value: Optional[BaseException],
+                 traceback: Optional[TracebackType]) -> Optional[bool]: ...
 
 class Mapping(_Collection[_KT], Generic[_KT, _VT_co]):
     # TODO: We wish the key type could also be covariant, but that doesn't work,

stdlib/3/unittest/__init__.pyi

@@ -3,12 +3,11 @@
 from typing import (
     Any, Callable, Dict, Iterable, Iterator, List, Optional, Pattern, Sequence,
     Set, FrozenSet, TextIO, Tuple, Type, TypeVar, Union, Generic,
-    overload,
+    overload, ContextManager
 )
 import logging
 import sys
 from types import ModuleType, TracebackType
-from contextlib import ContextManager
 
 
 _T = TypeVar('_T')

third_party/2and3/atomicwrites/__init__.pyi

@@ -1,16 +1,16 @@
-import contextlib
 import os
 import sys
 import tempfile
-from typing import Any, AnyStr, Callable, IO, Iterator, Text
+from typing import Any, AnyStr, Callable, ContextManager, IO, Iterator, Text
+
 def replace_atomic(src: AnyStr, dst: AnyStr) -> None: ...
 def move_atomic(src: AnyStr, dst: AnyStr) -> None: ...
 class AtomicWriter(object):
     def __init__(self, path: AnyStr, mode: Text='w', overwrite: bool=False) -> None: ...
-    def open(self) -> contextlib.ContextManager[IO]: ...
-    def _open(self, get_fileobject: Callable) -> contextlib.ContextManager[IO]: ...
+    def open(self) -> ContextManager[IO]: ...
+    def _open(self, get_fileobject: Callable) -> ContextManager[IO]: ...
     def get_fileobject(self, dir: AnyStr=None, **kwargs) -> IO: ...
     def sync(self, f: IO) -> None: ...
     def commit(self, f: IO) -> None: ...
     def rollback(self, f: IO) -> None: ...
-def atomic_write(path: AnyStr, writer_cls: type=AtomicWriter, **cls_kwargs) -> contextlib.ContextManager[IO]: ...
+def atomic_write(path: AnyStr, writer_cls: type=AtomicWriter, **cls_kwargs) -> ContextManager[IO]: ...