From 5734b9e5a9faa626f84d8b1d493d33f6b06669a9 Mon Sep 17 00:00:00 2001 From: sobolevn Date: Mon, 20 Dec 2021 00:52:12 +0300 Subject: [PATCH 1/5] Fixes generic inference in functions with `TypeGuard`, refs #11780 --- mypy/applytype.py | 7 +++++++ mypy/checkexpr.py | 21 ++++++++++++--------- test-data/unit/check-typeguard.test | 15 +++++++++++++++ 3 files changed, 34 insertions(+), 9 deletions(-) diff --git a/mypy/applytype.py b/mypy/applytype.py index 5b803a4aaa0b..e785246ece82 100644 --- a/mypy/applytype.py +++ b/mypy/applytype.py @@ -100,6 +100,12 @@ def apply_generic_arguments( # Apply arguments to argument types. arg_types = [expand_type(at, id_to_type) for at in callable.arg_types] + # Apply arguments to TypeGuard if any. + if callable.type_guard is not None: + type_guard = expand_type(callable.type_guard, id_to_type) + else: + type_guard = None + # The callable may retain some type vars if only some were applied. remaining_tvars = [tv for tv in tvars if tv.id not in id_to_type] @@ -107,4 +113,5 @@ def apply_generic_arguments( arg_types=arg_types, ret_type=expand_type(callable.ret_type, id_to_type), variables=remaining_tvars, + type_guard=type_guard, ) diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py index 1647339ef217..9240752d3641 100644 --- a/mypy/checkexpr.py +++ b/mypy/checkexpr.py @@ -336,11 +336,6 @@ def visit_call_expr_inner(self, e: CallExpr, allow_none_return: bool = False) -> ret_type=self.object_type(), fallback=self.named_type('builtins.function')) callee_type = get_proper_type(self.accept(e.callee, type_context, always_allow_any=True)) - if (isinstance(e.callee, RefExpr) - and isinstance(callee_type, CallableType) - and callee_type.type_guard is not None): - # Cache it for find_isinstance_check() - e.callee.type_guard = callee_type.type_guard if (self.chk.options.disallow_untyped_calls and self.chk.in_checked_function() and isinstance(callee_type, CallableType) @@ -882,10 +877,18 @@ def check_call_expr_with_callee_type(self, # Unions are special-cased to allow plugins to act on each item in the union. elif member is not None and isinstance(object_type, UnionType): return self.check_union_call_expr(e, object_type, member) - return self.check_call(callee_type, e.args, e.arg_kinds, e, - e.arg_names, callable_node=e.callee, - callable_name=callable_name, - object_type=object_type)[0] + ret_type, callee_type = self.check_call( + callee_type, e.args, e.arg_kinds, e, + e.arg_names, callable_node=e.callee, + callable_name=callable_name, + object_type=object_type, + ) + if (isinstance(e.callee, RefExpr) + and isinstance(callee_type, CallableType) + and callee_type.type_guard is not None): + # Cache it for find_isinstance_check() + e.callee.type_guard = callee_type.type_guard + return ret_type def check_union_call_expr(self, e: CallExpr, object_type: UnionType, member: str) -> Type: """"Type check calling a member expression where the base type is a union.""" diff --git a/test-data/unit/check-typeguard.test b/test-data/unit/check-typeguard.test index 32fe5e750989..603bf180188e 100644 --- a/test-data/unit/check-typeguard.test +++ b/test-data/unit/check-typeguard.test @@ -547,3 +547,18 @@ accepts_typeguard(with_typeguard_a) # E: Argument 1 to "accepts_typeguard" has accepts_typeguard(with_typeguard_b) accepts_typeguard(with_typeguard_c) [builtins fixtures/tuple.pyi] + +[case testTypeGuardWithGenerics] +from typing import TypeVar, Tuple +from typing_extensions import TypeGuard + +_T = TypeVar("_T") + +def is_two_element_tuple(val: Tuple[_T, ...]) -> TypeGuard[Tuple[_T, _T]]: + pass + +def func(names: Tuple[str, ...]): + reveal_type(names) # N: Revealed type is "builtins.tuple[builtins.str]" + if is_two_element_tuple(names): + reveal_type(names) # N: Revealed type is "Tuple[builtins.str*, builtins.str*]" +[builtins fixtures/tuple.pyi] From 9b5d2d8a666e57b7a5dda0035e4d28c41fa1973c Mon Sep 17 00:00:00 2001 From: sobolevn Date: Mon, 20 Dec 2021 00:58:02 +0300 Subject: [PATCH 2/5] Fixes generic inference in functions with `TypeGuard`, refs #11780 --- mypy/checkexpr.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py index 9240752d3641..0cabd490aba3 100644 --- a/mypy/checkexpr.py +++ b/mypy/checkexpr.py @@ -883,11 +883,12 @@ def check_call_expr_with_callee_type(self, callable_name=callable_name, object_type=object_type, ) + proper_callee = get_proper_type(callee_type) if (isinstance(e.callee, RefExpr) - and isinstance(callee_type, CallableType) - and callee_type.type_guard is not None): + and isinstance(proper_callee, CallableType) + and proper_callee.type_guard is not None): # Cache it for find_isinstance_check() - e.callee.type_guard = callee_type.type_guard + e.callee.type_guard = proper_callee.type_guard return ret_type def check_union_call_expr(self, e: CallExpr, object_type: UnionType, member: str) -> Type: From 4f14be3098a7dd2f323e3265da743ce7ee85899e Mon Sep 17 00:00:00 2001 From: sobolevn Date: Mon, 20 Dec 2021 11:28:13 +0300 Subject: [PATCH 3/5] More tests --- test-data/unit/check-typeguard.test | 37 ++++++++++++++++++++++++++++- 1 file changed, 36 insertions(+), 1 deletion(-) diff --git a/test-data/unit/check-typeguard.test b/test-data/unit/check-typeguard.test index 603bf180188e..8b444b60af5a 100644 --- a/test-data/unit/check-typeguard.test +++ b/test-data/unit/check-typeguard.test @@ -548,7 +548,42 @@ accepts_typeguard(with_typeguard_b) accepts_typeguard(with_typeguard_c) [builtins fixtures/tuple.pyi] -[case testTypeGuardWithGenerics] +[case testTypeGuardWithIdentityGeneric] +from typing import TypeVar +from typing_extensions import TypeGuard + +_T = TypeVar("_T") + +def identity(val: _T) -> TypeGuard[_T]: + pass + +def func1(name: _T): + reveal_type(name) # N: Revealed type is "_T`-1" + if identity(name): + reveal_type(name) # N: Revealed type is "_T`-1" + +def func2(name: str): + reveal_type(name) # N: Revealed type is "builtins.str" + if identity(name): + reveal_type(name) # N: Revealed type is "builtins.str*" +[builtins fixtures/tuple.pyi] + +[case testTypeGuardWithGenericInstance] +from typing import TypeVar, List +from typing_extensions import TypeGuard + +_T = TypeVar("_T") + +def is_list_of_str(val: _T) -> TypeGuard[List[_T]]: + pass + +def func(name: str): + reveal_type(name) # N: Revealed type is "builtins.str" + if is_list_of_str(name): + reveal_type(name) # N: Revealed type is "builtins.list[builtins.str*]" +[builtins fixtures/tuple.pyi] + +[case testTypeGuardWithTupleGeneric] from typing import TypeVar, Tuple from typing_extensions import TypeGuard From 3f280cb066b20ca3b2c9bd53c19b35a725d19ef1 Mon Sep 17 00:00:00 2001 From: Shantanu <12621235+hauntsaninja@users.noreply.github.com> Date: Sat, 7 May 2022 16:43:24 -0700 Subject: [PATCH 4/5] Apply suggestions from code review --- test-data/unit/check-typeguard.test | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/test-data/unit/check-typeguard.test b/test-data/unit/check-typeguard.test index 5bc59d79d167..58031469157c 100644 --- a/test-data/unit/check-typeguard.test +++ b/test-data/unit/check-typeguard.test @@ -565,7 +565,7 @@ def func1(name: _T): def func2(name: str): reveal_type(name) # N: Revealed type is "builtins.str" if identity(name): - reveal_type(name) # N: Revealed type is "builtins.str*" + reveal_type(name) # N: Revealed type is "builtins.str" [builtins fixtures/tuple.pyi] [case testTypeGuardWithGenericInstance] @@ -580,7 +580,7 @@ def is_list_of_str(val: _T) -> TypeGuard[List[_T]]: def func(name: str): reveal_type(name) # N: Revealed type is "builtins.str" if is_list_of_str(name): - reveal_type(name) # N: Revealed type is "builtins.list[builtins.str*]" + reveal_type(name) # N: Revealed type is "builtins.list[builtins.str]" [builtins fixtures/tuple.pyi] [case testTypeGuardWithTupleGeneric] @@ -595,5 +595,5 @@ def is_two_element_tuple(val: Tuple[_T, ...]) -> TypeGuard[Tuple[_T, _T]]: def func(names: Tuple[str, ...]): reveal_type(names) # N: Revealed type is "builtins.tuple[builtins.str]" if is_two_element_tuple(names): - reveal_type(names) # N: Revealed type is "Tuple[builtins.str*, builtins.str*]" + reveal_type(names) # N: Revealed type is "Tuple[builtins.str, builtins.str]" [builtins fixtures/tuple.pyi] From 7dbac9db7b30d004120fc20953a8468d5246e98d Mon Sep 17 00:00:00 2001 From: Shantanu <12621235+hauntsaninja@users.noreply.github.com> Date: Sat, 7 May 2022 17:00:51 -0700 Subject: [PATCH 5/5] Update test-data/unit/check-typeguard.test --- test-data/unit/check-typeguard.test | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test-data/unit/check-typeguard.test b/test-data/unit/check-typeguard.test index 58031469157c..64fc7ea695cb 100644 --- a/test-data/unit/check-typeguard.test +++ b/test-data/unit/check-typeguard.test @@ -593,7 +593,7 @@ def is_two_element_tuple(val: Tuple[_T, ...]) -> TypeGuard[Tuple[_T, _T]]: pass def func(names: Tuple[str, ...]): - reveal_type(names) # N: Revealed type is "builtins.tuple[builtins.str]" + reveal_type(names) # N: Revealed type is "builtins.tuple[builtins.str, ...]" if is_two_element_tuple(names): reveal_type(names) # N: Revealed type is "Tuple[builtins.str, builtins.str]" [builtins fixtures/tuple.pyi]