From 207200bbd134c4dc80791473c10411dbf3190c35 Mon Sep 17 00:00:00 2001 From: Barnaby Shearer Date: Mon, 18 Oct 2021 13:37:14 +0100 Subject: [PATCH 1/3] Extend contextmanager to asynccontextmanager Fixes #9922 --- mypy/plugins/default.py | 2 ++ test-data/unit/check-default-plugin.test | 18 ++++++++++++++++++ 2 files changed, 20 insertions(+) diff --git a/mypy/plugins/default.py b/mypy/plugins/default.py index 552a52c5c860..3890c64081d4 100644 --- a/mypy/plugins/default.py +++ b/mypy/plugins/default.py @@ -26,6 +26,8 @@ def get_function_hook(self, fullname: str if fullname == 'contextlib.contextmanager': return contextmanager_callback + if fullname == 'contextlib.asynccontextmanager': + return contextmanager_callback elif fullname == 'builtins.open' and self.python_version[0] == 3: return open_callback elif fullname == 'ctypes.Array': diff --git a/test-data/unit/check-default-plugin.test b/test-data/unit/check-default-plugin.test index 7493763d3e72..d4513044a473 100644 --- a/test-data/unit/check-default-plugin.test +++ b/test-data/unit/check-default-plugin.test @@ -24,6 +24,24 @@ f = g # E: Incompatible types in assignment (expression has type "Callable[[Any, [typing fixtures/typing-medium.pyi] [builtins fixtures/tuple.pyi] +[case testAsyncContextManagerWithGenericFunction] +# flags: --python-version 3.7 +from contextlib import asynccontextmanager +from typing import TypeVar, AsyncGenerator + +T = TypeVar('T') + +@asynccontextmanager +async def yield_id(item: T) -> AsyncGenerator[T, None]: + yield item + +reveal_type(yield_id) # N: Revealed type is "def [T] (item: T`-1) -> typing.AsyncContextManager[T`-1]" + +async with yield_id(1) as x: + reveal_type(x) # N: Revealed type is "builtins.int*" +[typing fixtures/typing-async.pyi] +[builtins fixtures/tuple.pyi] + [case testContextManagerWithUnspecifiedArguments] from contextlib import contextmanager from typing import Callable, Iterator From ab503543f5bc55bb103613672d0374f4e3f5bbeb Mon Sep 17 00:00:00 2001 From: Barnaby Shearer Date: Tue, 19 Oct 2021 10:45:44 +0100 Subject: [PATCH 2/3] PR-feedback: Combine ifs, SendTypeTest --- mypy/plugins/default.py | 4 +--- test-data/unit/check-default-plugin.test | 19 +++++++++++++++++++ 2 files changed, 20 insertions(+), 3 deletions(-) diff --git a/mypy/plugins/default.py b/mypy/plugins/default.py index 3890c64081d4..065703dffede 100644 --- a/mypy/plugins/default.py +++ b/mypy/plugins/default.py @@ -24,9 +24,7 @@ def get_function_hook(self, fullname: str ) -> Optional[Callable[[FunctionContext], Type]]: from mypy.plugins import ctypes - if fullname == 'contextlib.contextmanager': - return contextmanager_callback - if fullname == 'contextlib.asynccontextmanager': + if fullname in ('contextlib.contextmanager', 'contextlib.asynccontextmanager'): return contextmanager_callback elif fullname == 'builtins.open' and self.python_version[0] == 3: return open_callback diff --git a/test-data/unit/check-default-plugin.test b/test-data/unit/check-default-plugin.test index d4513044a473..239575ed41d7 100644 --- a/test-data/unit/check-default-plugin.test +++ b/test-data/unit/check-default-plugin.test @@ -42,6 +42,25 @@ async with yield_id(1) as x: [typing fixtures/typing-async.pyi] [builtins fixtures/tuple.pyi] +[case testAsyncContextManagerWithGenericFunctionAndSendType] +# flags: --python-version 3.7 +from contextlib import asynccontextmanager +from typing import TypeVar, AsyncGenerator + +T = TypeVar('T') +S = TypeVar('S') + +@asynccontextmanager +async def yield_id(item: T) -> AsyncGenerator[T, S]: + yield item + +reveal_type(yield_id) # N: Revealed type is "def [T, S] (item: T`-1) -> typing.AsyncContextManager[T`-1]" + +async with yield_id(1) as x: + reveal_type(x) # N: Revealed type is "builtins.int*" +[typing fixtures/typing-async.pyi] +[builtins fixtures/tuple.pyi] + [case testContextManagerWithUnspecifiedArguments] from contextlib import contextmanager from typing import Callable, Iterator From 7ecaedbd6e7427a73964c033adbd2aee9d5bedeb Mon Sep 17 00:00:00 2001 From: Barnaby Shearer Date: Tue, 19 Oct 2021 13:46:33 +0100 Subject: [PATCH 3/3] PR Feedback: Use detach_callable --- mypy/plugins/default.py | 5 +++-- test-data/unit/check-default-plugin.test | 28 +++++++++++++++++++++--- 2 files changed, 28 insertions(+), 5 deletions(-) diff --git a/mypy/plugins/default.py b/mypy/plugins/default.py index 065703dffede..633549fc3133 100644 --- a/mypy/plugins/default.py +++ b/mypy/plugins/default.py @@ -15,6 +15,7 @@ from mypy.subtypes import is_subtype from mypy.typeops import make_simplified_union from mypy.checkexpr import is_literal_type_like +from mypy.checker import detach_callable class DefaultPlugin(Plugin): @@ -183,12 +184,12 @@ def contextmanager_callback(ctx: FunctionContext) -> Type: and isinstance(default_return, CallableType)): # The stub signature doesn't preserve information about arguments so # add them back here. - return default_return.copy_modified( + return detach_callable(default_return.copy_modified( arg_types=arg_type.arg_types, arg_kinds=arg_type.arg_kinds, arg_names=arg_type.arg_names, variables=arg_type.variables, - is_ellipsis_args=arg_type.is_ellipsis_args) + is_ellipsis_args=arg_type.is_ellipsis_args)) return ctx.default_return_type diff --git a/test-data/unit/check-default-plugin.test b/test-data/unit/check-default-plugin.test index 239575ed41d7..fc9f132abb62 100644 --- a/test-data/unit/check-default-plugin.test +++ b/test-data/unit/check-default-plugin.test @@ -24,15 +24,37 @@ f = g # E: Incompatible types in assignment (expression has type "Callable[[Any, [typing fixtures/typing-medium.pyi] [builtins fixtures/tuple.pyi] +[case testContextManagerWithGenericFunctionAndSendType] +from contextlib import contextmanager +from typing import TypeVar, Generator + +T = TypeVar('T') +S = TypeVar('S') + +@contextmanager +def yield_id(item: T) -> Generator[T, S, None]: + yield item + +reveal_type(yield_id) # N: Revealed type is "def [T] (item: T`-1) -> contextlib.GeneratorContextManager[T`-1]" + +with yield_id(1) as x: + reveal_type(x) # N: Revealed type is "builtins.int*" + +f = yield_id +def g(x, y): pass +f = g # E: Incompatible types in assignment (expression has type "Callable[[Any, Any], Any]", variable has type "Callable[[T], GeneratorContextManager[T]]") +[typing fixtures/typing-medium.pyi] +[builtins fixtures/tuple.pyi] + [case testAsyncContextManagerWithGenericFunction] # flags: --python-version 3.7 from contextlib import asynccontextmanager -from typing import TypeVar, AsyncGenerator +from typing import TypeVar, AsyncIterator T = TypeVar('T') @asynccontextmanager -async def yield_id(item: T) -> AsyncGenerator[T, None]: +async def yield_id(item: T) -> AsyncIterator[T]: yield item reveal_type(yield_id) # N: Revealed type is "def [T] (item: T`-1) -> typing.AsyncContextManager[T`-1]" @@ -54,7 +76,7 @@ S = TypeVar('S') async def yield_id(item: T) -> AsyncGenerator[T, S]: yield item -reveal_type(yield_id) # N: Revealed type is "def [T, S] (item: T`-1) -> typing.AsyncContextManager[T`-1]" +reveal_type(yield_id) # N: Revealed type is "def [T] (item: T`-1) -> typing.AsyncContextManager[T`-1]" async with yield_id(1) as x: reveal_type(x) # N: Revealed type is "builtins.int*"