From 20d9fcf858777a17400a41d89891365f3a20fc26 Mon Sep 17 00:00:00 2001
From: nimin98 <nimin98@gmail.com>
Date: Tue, 11 Apr 2017 14:51:03 -0700
Subject: [PATCH] Support callable contextmanagers in contextlib (#1152)

mypy could not recognize the case when we use contextmanager as a
decorator, which has been supported since Python 3.2.

In the following code snippet,

from contextlib import contextmanager

@contextmanager
def foo(arg1):
    try:
        print(arg1)
        print('1')
        yield
    finally:
        print('2')

@foo('0')
def foo2():
    print('3')

foo2()

we get mypy error as follows,
error: ContextManager[Any] not callable

The suggested changes can fix this error and properly reflect the
updated contextmanager usage pattern.
---
 stdlib/2and3/contextlib.pyi | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/stdlib/2and3/contextlib.pyi b/stdlib/2and3/contextlib.pyi
index 0f14afaee..6ff18d846 100644
--- a/stdlib/2and3/contextlib.pyi
+++ b/stdlib/2and3/contextlib.pyi
@@ -20,7 +20,12 @@ class ContextManager(Generic[_T]):
                  exc_val: Optional[Exception],
                  exc_tb: Optional[TracebackType]) -> bool: ...
 
-def contextmanager(func: Callable[..., Iterator[_T]]) -> Callable[..., ContextManager[_T]]: ...
+if sys.version_info >= (3, 2):
+    class GeneratorContextManager(Generic[_T], ContextManager[_T]):
+        def __call__(self, func: Callable[..., _T]) -> Callable[..., _T]: ...
+    def contextmanager(func: Callable[..., Iterator[_T]]) -> Callable[..., GeneratorContextManager[_T]]: ...
+else:
+    def contextmanager(func: Callable[..., Iterator[_T]]) -> Callable[..., ContextManager[_T]]: ...
 
 if sys.version_info < (3,):
     def nested(*mgr: ContextManager[Any]) -> ContextManager[Iterable[Any]]: ...