From c8164dc8667702ed104e7064b231ec790571385b Mon Sep 17 00:00:00 2001 From: Ivan Levkivskyi Date: Sat, 19 Nov 2022 14:52:21 +0000 Subject: [PATCH 1/5] Try empty context when assigning to union typed variables --- mypy/checker.py | 17 ++++++++++ test-data/unit/check-inference-context.test | 37 +++++++++++++++++++++ test-data/unit/check-typeddict.test | 2 +- 3 files changed, 55 insertions(+), 1 deletion(-) diff --git a/mypy/checker.py b/mypy/checker.py index c7de4911501a..3c6f93c9a1c7 100644 --- a/mypy/checker.py +++ b/mypy/checker.py @@ -3870,6 +3870,23 @@ def check_simple_assignment( rvalue_type = self.expr_checker.accept( rvalue, lvalue_type, always_allow_any=always_allow_any ) + if isinstance(get_proper_type(lvalue_type), UnionType): + # Try re-inferring r.h.s. in empty context, and use that if it + # results in a narrower type. We don't do this always because this + # may cause some perf impact, plus we want to partially preserve + # the old behavior. This helps with various practical examples, see + # e.g. testOptionalTypeNarrowedByGenericCall. + with self.msg.filter_errors() as local_errors: + alt_rvalue_type = self.expr_checker.accept( + rvalue, None, always_allow_any=always_allow_any + ) + if ( + not local_errors.has_new_errors() + # Skip literal types, as they have special logic (for better errors). + and not isinstance(get_proper_type(rvalue_type), LiteralType) + and is_proper_subtype(alt_rvalue_type, rvalue_type) + ): + rvalue_type = alt_rvalue_type if isinstance(rvalue_type, DeletedType): self.msg.deleted_as_rvalue(rvalue_type, context) if isinstance(lvalue_type, DeletedType): diff --git a/test-data/unit/check-inference-context.test b/test-data/unit/check-inference-context.test index 2e26f54c6e93..71b620e54850 100644 --- a/test-data/unit/check-inference-context.test +++ b/test-data/unit/check-inference-context.test @@ -1419,3 +1419,40 @@ def bar(x: Union[Mapping[Any, Any], Dict[Any, Sequence[Any]]]) -> None: ... bar({1: 2}) [builtins fixtures/dict.pyi] + +[case testOptionalTypeNarrowedByGenericCall] +# flags: --strict-optional +from typing import Dict, Optional + +d: Dict[str, str] = {} + +def foo(arg: Optional[str] = None) -> None: + if arg is None: + arg = d.get("a", "b") + reveal_type(arg) # N: Revealed type is "builtins.str" +[builtins fixtures/dict.pyi] + +[case testOptionalTypeNarrowedByGenericCall2] +# flags: --strict-optional +from typing import Dict, Optional + +d: Dict[str, str] = {} +x: Optional[str] +if x: + reveal_type(x) # N: Revealed type is "builtins.str" + x = d.get(x, x) + reveal_type(x) # N: Revealed type is "builtins.str" +[builtins fixtures/dict.pyi] + +[case testOptionalTypeNarrowedByGenericCall3] +# flags: --strict-optional +from typing import Generic, TypeVar, Union + +T = TypeVar("T") +def bar(arg: Union[str, T]) -> Union[str, T]: ... + +def foo(arg: Union[str, int]) -> None: + if isinstance(arg, int): + arg = bar("default") + reveal_type(arg) # N: Revealed type is "builtins.str" +[builtins fixtures/isinstance.pyi] diff --git a/test-data/unit/check-typeddict.test b/test-data/unit/check-typeddict.test index 24521062a5d4..fbef6157087c 100644 --- a/test-data/unit/check-typeddict.test +++ b/test-data/unit/check-typeddict.test @@ -893,7 +893,7 @@ B = TypedDict('B', {'@type': Literal['b-type'], 'b': int}) c: Union[A, B] = {'@type': 'a-type', 'a': 'Test'} reveal_type(c) # N: Revealed type is "Union[TypedDict('__main__.A', {'@type': Literal['a-type'], 'a': builtins.str}), TypedDict('__main__.B', {'@type': Literal['b-type'], 'b': builtins.int})]" -[builtins fixtures/tuple.pyi] +[builtins fixtures/dict.pyi] [case testTypedDictUnionAmbiguousCase] from typing import Union, Mapping, Any, cast From 3c3b7594313652efa5ae83f17ecf811483c8cc1d Mon Sep 17 00:00:00 2001 From: Ivan Levkivskyi Date: Sat, 19 Nov 2022 17:38:15 +0000 Subject: [PATCH 2/5] Be more careful --- mypy/checker.py | 4 +++- test-data/unit/check-inference-context.test | 11 +++++++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/mypy/checker.py b/mypy/checker.py index 3c6f93c9a1c7..6f5d936cf91b 100644 --- a/mypy/checker.py +++ b/mypy/checker.py @@ -3876,7 +3876,7 @@ def check_simple_assignment( # may cause some perf impact, plus we want to partially preserve # the old behavior. This helps with various practical examples, see # e.g. testOptionalTypeNarrowedByGenericCall. - with self.msg.filter_errors() as local_errors: + with self.msg.filter_errors() as local_errors, self.local_type_map() as type_map: alt_rvalue_type = self.expr_checker.accept( rvalue, None, always_allow_any=always_allow_any ) @@ -3884,9 +3884,11 @@ def check_simple_assignment( not local_errors.has_new_errors() # Skip literal types, as they have special logic (for better errors). and not isinstance(get_proper_type(rvalue_type), LiteralType) + and is_valid_inferred_type(alt_rvalue_type) and is_proper_subtype(alt_rvalue_type, rvalue_type) ): rvalue_type = alt_rvalue_type + self.store_types(type_map) if isinstance(rvalue_type, DeletedType): self.msg.deleted_as_rvalue(rvalue_type, context) if isinstance(lvalue_type, DeletedType): diff --git a/test-data/unit/check-inference-context.test b/test-data/unit/check-inference-context.test index 71b620e54850..f8a96b45992c 100644 --- a/test-data/unit/check-inference-context.test +++ b/test-data/unit/check-inference-context.test @@ -1456,3 +1456,14 @@ def foo(arg: Union[str, int]) -> None: arg = bar("default") reveal_type(arg) # N: Revealed type is "builtins.str" [builtins fixtures/isinstance.pyi] + +[case testOptionalTypeNarrowedByGenericCall4] +# flags: --strict-optional +from typing import Optional, List, Generic, TypeVar + +T = TypeVar("T", covariant=True) +class C(Generic[T]): ... + +x: Optional[C[int]] = None +y = x = C() +reveal_type(y) # N: Revealed type is "__main__.C[builtins.int]" From 05f849c2ed499e3df186dc1a58afe3a7ed2aa4ac Mon Sep 17 00:00:00 2001 From: Ivan Levkivskyi Date: Sat, 19 Nov 2022 18:59:02 +0000 Subject: [PATCH 3/5] Handle one more edge case --- mypy/checker.py | 2 ++ test-data/unit/check-inference-context.test | 9 +++++++++ 2 files changed, 11 insertions(+) diff --git a/mypy/checker.py b/mypy/checker.py index 6f5d936cf91b..6ac4c0968939 100644 --- a/mypy/checker.py +++ b/mypy/checker.py @@ -3884,6 +3884,8 @@ def check_simple_assignment( not local_errors.has_new_errors() # Skip literal types, as they have special logic (for better errors). and not isinstance(get_proper_type(rvalue_type), LiteralType) + # Skip Any type, since it is special cased in binder. + and not isinstance(get_proper_type(alt_rvalue_type), AnyType) and is_valid_inferred_type(alt_rvalue_type) and is_proper_subtype(alt_rvalue_type, rvalue_type) ): diff --git a/test-data/unit/check-inference-context.test b/test-data/unit/check-inference-context.test index f8a96b45992c..f80f93eb2615 100644 --- a/test-data/unit/check-inference-context.test +++ b/test-data/unit/check-inference-context.test @@ -1467,3 +1467,12 @@ class C(Generic[T]): ... x: Optional[C[int]] = None y = x = C() reveal_type(y) # N: Revealed type is "__main__.C[builtins.int]" + +[case testOptionalTypeNarrowedByGenericCall5] +from typing import Any, Tuple, Union + +i: Union[Tuple[Any, ...], int] +b: Any +i = i if isinstance(i, int) else b +reveal_type(i) # N: Revealed type is "Union[Any, builtins.int]" +[builtins fixtures/isinstance.pyi] From a6c67d32a2675cd6e0a1bcf4b71c2dc3fd547f72 Mon Sep 17 00:00:00 2001 From: Ivan Levkivskyi Date: Mon, 21 Nov 2022 15:11:35 +0000 Subject: [PATCH 4/5] Address CR (perf) --- mypy/checker.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/mypy/checker.py b/mypy/checker.py index 6ac4c0968939..21c48f1c9ef1 100644 --- a/mypy/checker.py +++ b/mypy/checker.py @@ -3870,7 +3870,11 @@ def check_simple_assignment( rvalue_type = self.expr_checker.accept( rvalue, lvalue_type, always_allow_any=always_allow_any ) - if isinstance(get_proper_type(lvalue_type), UnionType): + if ( + isinstance(get_proper_type(lvalue_type), UnionType) + # Skip literal types, as they have special logic (for better errors). + and not isinstance(get_proper_type(rvalue_type), LiteralType) + ): # Try re-inferring r.h.s. in empty context, and use that if it # results in a narrower type. We don't do this always because this # may cause some perf impact, plus we want to partially preserve @@ -3882,8 +3886,6 @@ def check_simple_assignment( ) if ( not local_errors.has_new_errors() - # Skip literal types, as they have special logic (for better errors). - and not isinstance(get_proper_type(rvalue_type), LiteralType) # Skip Any type, since it is special cased in binder. and not isinstance(get_proper_type(alt_rvalue_type), AnyType) and is_valid_inferred_type(alt_rvalue_type) From 233d0dcbbea7995f48f4669efecb6d115604937a Mon Sep 17 00:00:00 2001 From: Ivan Levkivskyi Date: Mon, 21 Nov 2022 19:28:06 +0000 Subject: [PATCH 5/5] Some perf optimization --- mypy/checker.py | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/mypy/checker.py b/mypy/checker.py index 21c48f1c9ef1..c21cf82f1c43 100644 --- a/mypy/checker.py +++ b/mypy/checker.py @@ -76,6 +76,7 @@ AssignmentStmt, Block, BreakStmt, + BytesExpr, CallExpr, ClassDef, ComparisonExpr, @@ -86,6 +87,7 @@ EllipsisExpr, Expression, ExpressionStmt, + FloatExpr, ForStmt, FuncBase, FuncDef, @@ -115,6 +117,7 @@ ReturnStmt, StarExpr, Statement, + StrExpr, SymbolNode, SymbolTable, SymbolTableNode, @@ -3849,6 +3852,23 @@ def inference_error_fallback_type(self, type: Type) -> Type: # we therefore need to erase them. return erase_typevars(fallback) + def simple_rvalue(self, rvalue: Expression) -> bool: + """Returns True for expressions for which inferred type should not depend on context. + + Note that this function can still return False for some expressions where inferred type + does not depend on context. It only exists for performance optimizations. + """ + if isinstance(rvalue, (IntExpr, StrExpr, BytesExpr, FloatExpr, RefExpr)): + return True + if isinstance(rvalue, CallExpr): + if isinstance(rvalue.callee, RefExpr) and isinstance(rvalue.callee.node, FuncBase): + typ = rvalue.callee.node.type + if isinstance(typ, CallableType): + return not typ.variables + elif isinstance(typ, Overloaded): + return not any(item.variables for item in typ.items) + return False + def check_simple_assignment( self, lvalue_type: Type | None, @@ -3874,6 +3894,7 @@ def check_simple_assignment( isinstance(get_proper_type(lvalue_type), UnionType) # Skip literal types, as they have special logic (for better errors). and not isinstance(get_proper_type(rvalue_type), LiteralType) + and not self.simple_rvalue(rvalue) ): # Try re-inferring r.h.s. in empty context, and use that if it # results in a narrower type. We don't do this always because this