Skip to content

Commit

Permalink
Fix regression when a contextmanager yields a generic (#11870)
Browse files Browse the repository at this point in the history
Resolves #11852
  • Loading branch information
asottile authored Jan 5, 2022
1 parent 9c05d3d commit 06afa7c
Show file tree
Hide file tree
Showing 2 changed files with 33 additions and 45 deletions.
5 changes: 2 additions & 3 deletions mypy/plugins/default.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,6 @@
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):
Expand Down Expand Up @@ -192,12 +191,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 detach_callable(default_return.copy_modified(
return 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


Expand Down
73 changes: 31 additions & 42 deletions test-data/unit/check-default-plugin.test
Original file line number Diff line number Diff line change
Expand Up @@ -24,56 +24,15 @@ 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, AsyncIterator

T = TypeVar('T')

@asynccontextmanager
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]"

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 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]:
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]"
Expand All @@ -83,6 +42,36 @@ async with yield_id(1) as x:
[typing fixtures/typing-async.pyi]
[builtins fixtures/tuple.pyi]

[case testContextManagerReturnsGenericFunction]
import contextlib
from typing import Callable
from typing import Generator
from typing import Iterable
from typing import TypeVar

TArg = TypeVar('TArg')
TRet = TypeVar('TRet')

@contextlib.contextmanager
def _thread_mapper(maxsize: int) -> Generator[
Callable[[Callable[[TArg], TRet], Iterable[TArg]], Iterable[TRet]],
None, None,
]:
# defined inline as there isn't a builtins.map fixture
def my_map(f: Callable[[TArg], TRet], it: Iterable[TArg]) -> Iterable[TRet]:
for x in it:
yield f(x)

yield my_map

def identity(x: int) -> int: return x

with _thread_mapper(1) as m:
lst = list(m(identity, [2, 3]))
reveal_type(lst) # N: Revealed type is "builtins.list[builtins.int*]"
[typing fixtures/typing-medium.pyi]
[builtins fixtures/list.pyi]

[case testContextManagerWithUnspecifiedArguments]
from contextlib import contextmanager
from typing import Callable, Iterator
Expand Down

0 comments on commit 06afa7c

Please sign in to comment.