From eecc3b3d9ceda204307354d272e9f3adb26a1fc2 Mon Sep 17 00:00:00 2001 From: "Michael J. Sullivan" Date: Fri, 20 Aug 2021 11:58:02 -0700 Subject: [PATCH 1/7] Support making a variable Optional in an else branch That is, support patterns such as: ``` if condition: foo = Foo() else: foo = None ``` Currently this does not work, but the the *reverse* does (because foo will be inferred as a PartialType). I think it might be worth tackling this in a more general way, for other types, though I think that is a little fiddlier and likely to be more controversial, so I'm starting with something special-cased for the "assigning literal None" case first. The rule we implement is that we allow updating the type of a variable when assigning `None` to it if the variable's type was inferred and it was defined in an earlier branch of the same `if/then/else` statement. Some infrastructure is added to make determinations about that. Given that this is probably my single biggest frustration with mypy, and that this PR took me less than three hours to prepare, I am pretty angry at myself for not having done this three years ago. --- mypy/binder.py | 19 +++++-- mypy/checker.py | 40 ++++++++++++- test-data/unit/check-optional.test | 90 ++++++++++++++++++++++++++++++ 3 files changed, 142 insertions(+), 7 deletions(-) diff --git a/mypy/binder.py b/mypy/binder.py index 1c711ce9c6318..34f39b81ee44e 100644 --- a/mypy/binder.py +++ b/mypy/binder.py @@ -31,9 +31,11 @@ class Frame: that were assigned in that frame. """ - def __init__(self) -> None: + def __init__(self, id: int) -> None: + self.id = id self.types: Dict[Key, Type] = {} self.unreachable = False + self.outer_if_frame = False # Should be set only if we're entering a frame where it's not # possible to accurately determine whether or not contained @@ -72,13 +74,15 @@ class A: type_assignments: Optional[Assigns] = None def __init__(self) -> None: + self.next_id = 1 + # The stack of frames currently used. These map # literal_hash(expr) -- literals like 'foo.bar' -- # to types. The last element of this list is the # top-most, current frame. Each earlier element # records the state as of when that frame was last # on top of the stack. - self.frames = [Frame()] + self.frames = [Frame(self._get_id())] # For frames higher in the stack, we record the set of # Frames that can escape there, either by falling off @@ -101,6 +105,10 @@ def __init__(self) -> None: self.break_frames: List[int] = [] self.continue_frames: List[int] = [] + def _get_id(self) -> int: + self.next_id += 1 + return self.next_id + def _add_dependencies(self, key: Key, value: Optional[Key] = None) -> None: if value is None: value = key @@ -111,7 +119,7 @@ def _add_dependencies(self, key: Key, value: Optional[Key] = None) -> None: def push_frame(self) -> Frame: """Push a new frame into the binder.""" - f = Frame() + f = Frame(self._get_id()) self.frames.append(f) self.options_on_return.append([]) return f @@ -349,7 +357,7 @@ def allow_jump(self, index: int) -> None: # so make sure the index is positive if index < 0: index += len(self.options_on_return) - frame = Frame() + frame = Frame(self._get_id()) for f in self.frames[index + 1:]: frame.types.update(f.types) if f.unreachable: @@ -367,6 +375,7 @@ def handle_continue(self) -> None: @contextmanager def frame_context(self, *, can_skip: bool, fall_through: int = 1, break_frame: int = 0, continue_frame: int = 0, + outer_if_frame: bool = False, try_frame: bool = False) -> Iterator[Frame]: """Return a context manager that pushes/pops frames on enter/exit. @@ -402,6 +411,8 @@ def frame_context(self, *, can_skip: bool, fall_through: int = 1, self.try_frames.add(len(self.frames) - 1) new_frame = self.push_frame() + if outer_if_frame: + new_frame.outer_if_frame = True if try_frame: # An exception may occur immediately self.allow_jump(-1) diff --git a/mypy/checker.py b/mypy/checker.py index ba020f5d97d5b..f627eb0b7a257 100644 --- a/mypy/checker.py +++ b/mypy/checker.py @@ -22,14 +22,14 @@ Context, Decorator, PrintStmt, BreakStmt, PassStmt, ContinueStmt, ComparisonExpr, StarExpr, EllipsisExpr, RefExpr, PromoteExpr, Import, ImportFrom, ImportAll, ImportBase, TypeAlias, - ARG_POS, ARG_STAR, LITERAL_TYPE, MDEF, GDEF, + ARG_POS, ARG_STAR, LITERAL_TYPE, LDEF, MDEF, GDEF, CONTRAVARIANT, COVARIANT, INVARIANT, TypeVarExpr, AssignmentExpr, is_final_node, ARG_NAMED) from mypy import nodes from mypy import operators from mypy.literals import literal, literal_hash, Key -from mypy.typeanal import has_any_from_unimported_type, check_for_explicit_any +from mypy.typeanal import has_any_from_unimported_type, check_for_explicit_any, make_optional_type from mypy.types import ( Type, AnyType, CallableType, FunctionLike, Overloaded, TupleType, TypedDictType, Instance, NoneType, strip_type, TypeType, TypeOfAny, @@ -203,6 +203,12 @@ class TypeChecker(NodeVisitor[None], CheckerPluginInterface): # directly or indirectly. module_refs: Set[str] + # A map from variable nodes to a snapshot of the current frames + # and whether they are associated with `if` statements. This can + # be used to determine if a variable is defined in a different + # branch of the same `if` statement. + var_decl_frames: Dict[Var, Dict[int, bool]] + # Plugin that provides special type checking rules for specific library # functions such as open(), etc. plugin: Plugin @@ -229,6 +235,7 @@ def __init__(self, errors: Errors, modules: Dict[str, MypyFile], options: Option self.dynamic_funcs = [] self.partial_types = [] self.partial_reported = set() + self.var_decl_frames = {} self.deferred_nodes = [] self.type_map = {} self.module_refs = set() @@ -2167,6 +2174,30 @@ def check_assignment(self, lvalue: Lvalue, rvalue: Expression, infer_lvalue_type rvalue_type, lvalue_type, infer_lvalue_type = self.check_member_assignment( instance_type, lvalue_type, rvalue, context=rvalue) else: + # Hacky special case for assigning a literal None + # to a variable defined in a previous if + # branch. When we detect this, we'll go back and + # make the type optional. This is somewhat + # unpleasant, and a generalization of this would + # be an improvement! + if (is_literal_none(rvalue) and + isinstance(lvalue, NameExpr) and + lvalue.kind == LDEF and + isinstance(lvalue.node, Var) and + lvalue.node.type and + lvalue.node in self.var_decl_frames): + decl_frame_map = self.var_decl_frames[lvalue.node] + # Check if the nearest common ancestor frame for the definition site + # and the current site is the enclosing frame of an if/elif/else block. + has_if_ancestor = False + for frame in reversed(self.binder.frames): + if frame.id in decl_frame_map: + has_if_ancestor = decl_frame_map[frame.id] + break + if has_if_ancestor: + lvalue_type = make_optional_type(lvalue_type) + self.set_inferred_type(lvalue.node, lvalue, lvalue_type) + rvalue_type = self.check_simple_assignment(lvalue_type, rvalue, context=rvalue, code=codes.ASSIGNMENT) @@ -2992,6 +3023,9 @@ def set_inferred_type(self, var: Var, lvalue: Lvalue, type: Type) -> None: if var and not self.current_node_deferred: var.type = type var.is_inferred = True + if var not in self.var_decl_frames: + self.var_decl_frames[var] = { + frame.id: frame.outer_if_frame for frame in self.binder.frames} if isinstance(lvalue, MemberExpr) and self.inferred_attribute_types is not None: # Store inferred attribute type so that we can check consistency afterwards. if lvalue.def_var is not None: @@ -3298,7 +3332,7 @@ def visit_if_stmt(self, s: IfStmt) -> None: """Type check an if statement.""" # This frame records the knowledge from previous if/elif clauses not being taken. # Fall-through to the original frame is handled explicitly in each block. - with self.binder.frame_context(can_skip=False, fall_through=0): + with self.binder.frame_context(can_skip=False, outer_if_frame=True, fall_through=0): for e, b in zip(s.expr, s.body): t = get_proper_type(self.expr_checker.accept(e)) diff --git a/test-data/unit/check-optional.test b/test-data/unit/check-optional.test index eedd5899c9497..60b51a3755b2b 100644 --- a/test-data/unit/check-optional.test +++ b/test-data/unit/check-optional.test @@ -832,3 +832,93 @@ main:4: error: Argument 1 to "asdf" has incompatible type "List[str]"; expected main:4: note: "List" is invariant -- see https://mypy.readthedocs.io/en/stable/common_issues.html#variance main:4: note: Consider using "Sequence" instead, which is covariant [builtins fixtures/list.pyi] + +[case testOptionalBackwards] +from typing import Optional + +def f1(b: bool) -> Optional[int]: + if b: + z = 10 + reveal_type(z) # N: Revealed type is "builtins.int" + else: + z = None + reveal_type(z) # N: Revealed type is "None" + reveal_type(z) # N: Revealed type is "Union[builtins.int, None]" + return z + +def f2(b: bool) -> int: + if b: + z = 10 + else: + z = None + return z # E: Incompatible return value type (got "Optional[int]", expected "int") + +def f3(b: bool) -> int: + # XXX: This one is a little questionable! Maybe we *do* want to allow this? + z = 10 + if b: + z = None # E: Incompatible types in assignment (expression has type "None", variable has type "int") + return z + +def f4() -> Optional[int]: + z = 10 + z = None # E: Incompatible types in assignment (expression has type "None", variable has type "int") + + return z + +def f5() -> None: + z = 10 + + def f() -> None: + nonlocal z + z = None # E: Incompatible types in assignment (expression has type "None", variable has type "int") + +def f6(b: bool) -> None: + if b: + z = 10 + else: + z = 11 + + def f() -> None: + nonlocal z + z = None # E: Incompatible types in assignment (expression has type "None", variable has type "int") + +def f7(b: bool) -> None: + if b: + z = 10 + else: + z = 11 + + z = None # E: Incompatible types in assignment (expression has type "None", variable has type "int") + +def f8(b: bool, c: bool) -> Optional[int]: + if b: + if c: + z = 10 + else: + z = 11 + else: + z = None + return z + +def f9(b: bool) -> None: + if b: + z: int = 10 + else: + z = None # E: Incompatible types in assignment (expression has type "None", variable has type "int") + +def f10(b: bool) -> None: + z: int + if b: + z = 10 + else: + z = None # E: Incompatible types in assignment (expression has type "None", variable has type "int") + + +def f11(b: bool, c: bool) -> None: + if b: + z = 10 + elif c: + z = 30 + else: + z = None From bfd3727a9aa57da34bd36826e22464c2c06f80ba Mon Sep 17 00:00:00 2001 From: "Michael J. Sullivan" Date: Fri, 20 Aug 2021 13:46:49 -0700 Subject: [PATCH 2/7] Don't update the type if it is Any --- mypy/checker.py | 3 ++- test-data/unit/check-optional.test | 10 ++++++++-- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/mypy/checker.py b/mypy/checker.py index f627eb0b7a257..180fb04af11a1 100644 --- a/mypy/checker.py +++ b/mypy/checker.py @@ -2185,7 +2185,8 @@ def check_assignment(self, lvalue: Lvalue, rvalue: Expression, infer_lvalue_type lvalue.kind == LDEF and isinstance(lvalue.node, Var) and lvalue.node.type and - lvalue.node in self.var_decl_frames): + lvalue.node in self.var_decl_frames and + not isinstance(get_proper_type(lvalue_type), AnyType)): decl_frame_map = self.var_decl_frames[lvalue.node] # Check if the nearest common ancestor frame for the definition site # and the current site is the enclosing frame of an if/elif/else block. diff --git a/test-data/unit/check-optional.test b/test-data/unit/check-optional.test index 60b51a3755b2b..2f648a1fad90b 100644 --- a/test-data/unit/check-optional.test +++ b/test-data/unit/check-optional.test @@ -834,7 +834,7 @@ main:4: note: Consider using "Sequence" instead, which is covariant [builtins fixtures/list.pyi] [case testOptionalBackwards] -from typing import Optional +from typing import Any, Optional def f1(b: bool) -> Optional[int]: if b: @@ -914,7 +914,6 @@ def f10(b: bool) -> None: else: z = None # E: Incompatible types in assignment (expression has type "None", variable has type "int") - def f11(b: bool, c: bool) -> None: if b: z = 10 @@ -922,3 +921,10 @@ def f11(b: bool, c: bool) -> None: z = 30 else: z = None + +def f12(b: bool, a: Any) -> None: + if b: + z = a + else: + z = None + reveal_type(z) # N: Revealed type is "Any" From 5859a7927dfddd449743c893aa9ba93adbb628e3 Mon Sep 17 00:00:00 2001 From: "Michael J. Sullivan" Date: Fri, 20 Aug 2021 16:07:22 -0700 Subject: [PATCH 3/7] thread outer_if_frame in and rename it --- mypy/binder.py | 14 ++++++-------- mypy/checker.py | 4 ++-- 2 files changed, 8 insertions(+), 10 deletions(-) diff --git a/mypy/binder.py b/mypy/binder.py index 34f39b81ee44e..3c0c4e793359a 100644 --- a/mypy/binder.py +++ b/mypy/binder.py @@ -31,11 +31,11 @@ class Frame: that were assigned in that frame. """ - def __init__(self, id: int) -> None: + def __init__(self, id: int, conditional_frame: bool = False) -> None: self.id = id self.types: Dict[Key, Type] = {} self.unreachable = False - self.outer_if_frame = False + self.conditional_frame = conditional_frame # Should be set only if we're entering a frame where it's not # possible to accurately determine whether or not contained @@ -117,9 +117,9 @@ def _add_dependencies(self, key: Key, value: Optional[Key] = None) -> None: for elt in subkeys(key): self._add_dependencies(elt, value) - def push_frame(self) -> Frame: + def push_frame(self, conditional_frame: bool = False) -> Frame: """Push a new frame into the binder.""" - f = Frame(self._get_id()) + f = Frame(self._get_id(), conditional_frame) self.frames.append(f) self.options_on_return.append([]) return f @@ -375,7 +375,7 @@ def handle_continue(self) -> None: @contextmanager def frame_context(self, *, can_skip: bool, fall_through: int = 1, break_frame: int = 0, continue_frame: int = 0, - outer_if_frame: bool = False, + conditional_frame: bool = False, try_frame: bool = False) -> Iterator[Frame]: """Return a context manager that pushes/pops frames on enter/exit. @@ -410,9 +410,7 @@ def frame_context(self, *, can_skip: bool, fall_through: int = 1, if try_frame: self.try_frames.add(len(self.frames) - 1) - new_frame = self.push_frame() - if outer_if_frame: - new_frame.outer_if_frame = True + new_frame = self.push_frame(conditional_frame) if try_frame: # An exception may occur immediately self.allow_jump(-1) diff --git a/mypy/checker.py b/mypy/checker.py index 180fb04af11a1..2c55867a92bfd 100644 --- a/mypy/checker.py +++ b/mypy/checker.py @@ -3026,7 +3026,7 @@ def set_inferred_type(self, var: Var, lvalue: Lvalue, type: Type) -> None: var.is_inferred = True if var not in self.var_decl_frames: self.var_decl_frames[var] = { - frame.id: frame.outer_if_frame for frame in self.binder.frames} + frame.id: frame.conditional_frame for frame in self.binder.frames} if isinstance(lvalue, MemberExpr) and self.inferred_attribute_types is not None: # Store inferred attribute type so that we can check consistency afterwards. if lvalue.def_var is not None: @@ -3333,7 +3333,7 @@ def visit_if_stmt(self, s: IfStmt) -> None: """Type check an if statement.""" # This frame records the knowledge from previous if/elif clauses not being taken. # Fall-through to the original frame is handled explicitly in each block. - with self.binder.frame_context(can_skip=False, outer_if_frame=True, fall_through=0): + with self.binder.frame_context(can_skip=False, conditional_frame=True, fall_through=0): for e, b in zip(s.expr, s.body): t = get_proper_type(self.expr_checker.accept(e)) From 08f09eafb468add0421bab8a6c9faa9c2130c7c6 Mon Sep 17 00:00:00 2001 From: "Michael J. Sullivan" Date: Fri, 20 Aug 2021 16:15:13 -0700 Subject: [PATCH 4/7] sure, some more tests --- test-data/unit/check-optional.test | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/test-data/unit/check-optional.test b/test-data/unit/check-optional.test index 2f648a1fad90b..bc90ef7494a24 100644 --- a/test-data/unit/check-optional.test +++ b/test-data/unit/check-optional.test @@ -928,3 +928,21 @@ def f12(b: bool, a: Any) -> None: else: z = None reveal_type(z) # N: Revealed type is "Any" + +def f13(b: bool, a: Any) -> None: + if b: + try: + z = f2(True) + except Exception: + raise RuntimeError + else: + z = None + +def f14(b: bool, a: Any) -> None: + if b: + with a: + z = 10 + else: + z = None + +[builtins fixtures/exception.pyi] From 2b24306598a912f63fdbcf543f517585ec063e00 Mon Sep 17 00:00:00 2001 From: "Michael J. Sullivan" Date: Fri, 20 Aug 2021 16:26:02 -0700 Subject: [PATCH 5/7] try it with try/loops --- mypy/checker.py | 4 ++-- test-data/unit/check-optional.test | 16 ++++++++++++++++ 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/mypy/checker.py b/mypy/checker.py index 2c55867a92bfd..73c81039ea4d2 100644 --- a/mypy/checker.py +++ b/mypy/checker.py @@ -418,7 +418,7 @@ def accept_loop(self, body: Statement, else_body: Optional[Statement] = None, *, Then check the else_body. """ # The outer frame accumulates the results of all iterations - with self.binder.frame_context(can_skip=False): + with self.binder.frame_context(can_skip=False, conditional_frame=True): while True: with self.binder.frame_context(can_skip=True, break_frame=2, continue_frame=1): @@ -3472,7 +3472,7 @@ def visit_try_without_finally(self, s: TryStmt, try_frame: bool) -> None: # was the top frame on entry. with self.binder.frame_context(can_skip=False, fall_through=2, try_frame=try_frame): # This frame receives exit via exception, and runs exception handlers - with self.binder.frame_context(can_skip=False, fall_through=2): + with self.binder.frame_context(can_skip=False, conditional_frame=True, fall_through=2): # Finally, the body of the try statement with self.binder.frame_context(can_skip=False, fall_through=2, try_frame=True): self.accept(s.body) diff --git a/test-data/unit/check-optional.test b/test-data/unit/check-optional.test index bc90ef7494a24..7a67052c45903 100644 --- a/test-data/unit/check-optional.test +++ b/test-data/unit/check-optional.test @@ -945,4 +945,20 @@ def f14(b: bool, a: Any) -> None: else: z = None +def f15() -> None: + try: + z = f2(True) + except Exception: + z = None + reveal_type(z) # N: Revealed type is "Union[builtins.int, None]" + +def f16(z: Any) -> None: + for x in z: + if x == 0: + y = 50 + break + else: + y = None + reveal_type(y) # N: Revealed type is "Union[builtins.int, None]" + [builtins fixtures/exception.pyi] From 481756027ffb4ea7db4362f2820bdbcf7c1ee1fb Mon Sep 17 00:00:00 2001 From: "Michael J. Sullivan" Date: Wed, 25 Aug 2021 07:57:21 -0700 Subject: [PATCH 6/7] Implement most of the tests that Jukka requested --- mypy/checker.py | 1 + test-data/unit/check-optional.test | 94 +++++++++++++++++++++++++++++- 2 files changed, 94 insertions(+), 1 deletion(-) diff --git a/mypy/checker.py b/mypy/checker.py index 73c81039ea4d2..c8d6ea0cd803b 100644 --- a/mypy/checker.py +++ b/mypy/checker.py @@ -4955,6 +4955,7 @@ def handle_partial_var_type( if is_local or not self.options.allow_untyped_globals: self.msg.need_annotation_for_var(node, context, self.options.python_version) + self.partial_reported.add(node) else: # Defer the node -- we might get a better type in the outer scope self.handle_cannot_determine_type(node.name, context) diff --git a/test-data/unit/check-optional.test b/test-data/unit/check-optional.test index 7a67052c45903..0ff7c2bc5b724 100644 --- a/test-data/unit/check-optional.test +++ b/test-data/unit/check-optional.test @@ -833,7 +833,7 @@ main:4: note: "List" is invariant -- see https://mypy.readthedocs.io/en/stable/c main:4: note: Consider using "Sequence" instead, which is covariant [builtins fixtures/list.pyi] -[case testOptionalBackwards] +[case testOptionalBackwards1] from typing import Any, Optional def f1(b: bool) -> Optional[int]: @@ -961,4 +961,96 @@ def f16(z: Any) -> None: y = None reveal_type(y) # N: Revealed type is "Union[builtins.int, None]" +def f17(b: bool, c: bool, d: bool) -> None: + if b: + z = 2 + elif c: + z = None + elif d: + z = 3 + reveal_type(z) # N: Revealed type is "Union[builtins.int, None]" + +def f18(b: bool, c: bool, d: bool) -> None: + if b: + z = 4 + else: + if c: + z = 5 + else: + z = None + reveal_type(z) # N: Revealed type is "Union[builtins.int, None]" + +def f19(b: bool, c: bool, d: bool) -> None: + if b: + z = 5 + else: + z = None + if c: + z = 6 + reveal_type(z) # N: Revealed type is "Union[builtins.int, None]" + +def f20(b: bool) -> None: + if b: + x: Any = 5 + else: + x = None + reveal_type(x) # N: Revealed type is "Any" + [builtins fixtures/exception.pyi] + +[case testOptionalBackwards2] + +def f1(b: bool) -> None: + if b: + x = [] # E: Need type annotation for "x" (hint: "x: List[] = ...") + else: + x = None + +def f2(b: bool) -> None: + if b: + x = [] + x.append(1) + else: + x = None + reveal_type(x) # N: Revealed type is "Union[builtins.list[builtins.int], None]" + + +[builtins fixtures/list.pyi] + +[case testOptionalBackwards3] + +# We don't allow this sort of updating for globals or attributes currently. +gb: bool +if gb: + Z = 10 +else: + Z = None # E: Incompatible types in assignment (expression has type "None", variable has type "int") + +class Foo: + def __init__(self, b: bool) -> None: + if b: + self.x = 5 + else: + self.x = None # E: Incompatible types in assignment (expression has type "None", variable has type "int") + + def foo(self) -> None: + reveal_type(self.x) # N: Revealed type is "builtins.int" + +[case testOptionalBackwards4] +from typing import Any, Optional + +def f1(b: bool) -> Optional[int]: + if b: + z = 10 + reveal_type(z) # N: Revealed type is "builtins.int" + else: + # Force the node to get deferred between the two assignments + Defer().defer + z = None + reveal_type(z) # N: Revealed type is "None" + reveal_type(z) # N: Revealed type is "Union[builtins.int, None]" + return z + +class Defer: + def __init__(self) -> None: + self.defer = 10 From 804f52e5698f425490f873fe16b6c054f7497a11 Mon Sep 17 00:00:00 2001 From: "Michael J. Sullivan" Date: Thu, 26 Aug 2021 16:33:08 -0700 Subject: [PATCH 7/7] tweaks --- mypy/checker.py | 18 ++++++++++-------- test-data/unit/check-optional.test | 23 +++++++++++++++++++++++ 2 files changed, 33 insertions(+), 8 deletions(-) diff --git a/mypy/checker.py b/mypy/checker.py index c8d6ea0cd803b..a6b732dc025ec 100644 --- a/mypy/checker.py +++ b/mypy/checker.py @@ -203,11 +203,13 @@ class TypeChecker(NodeVisitor[None], CheckerPluginInterface): # directly or indirectly. module_refs: Set[str] - # A map from variable nodes to a snapshot of the current frames - # and whether they are associated with `if` statements. This can - # be used to determine if a variable is defined in a different - # branch of the same `if` statement. - var_decl_frames: Dict[Var, Dict[int, bool]] + # A map from variable nodes to a snapshot of the frame ids of the + # frames that were active when the variable was declared. This can + # be used to determine nearest common ancestor frame of a variable's + # declaration and the current frame, which lets us determine if it + # was declared in a different branch of the same `if` statement + # (if that frame is a conditional_frame). + var_decl_frames: Dict[Var, Set[int]] # Plugin that provides special type checking rules for specific library # functions such as open(), etc. @@ -2193,7 +2195,7 @@ def check_assignment(self, lvalue: Lvalue, rvalue: Expression, infer_lvalue_type has_if_ancestor = False for frame in reversed(self.binder.frames): if frame.id in decl_frame_map: - has_if_ancestor = decl_frame_map[frame.id] + has_if_ancestor = frame.conditional_frame break if has_if_ancestor: lvalue_type = make_optional_type(lvalue_type) @@ -3025,8 +3027,8 @@ def set_inferred_type(self, var: Var, lvalue: Lvalue, type: Type) -> None: var.type = type var.is_inferred = True if var not in self.var_decl_frames: - self.var_decl_frames[var] = { - frame.id: frame.conditional_frame for frame in self.binder.frames} + # Used for the hack to improve optional type inference in conditionals + self.var_decl_frames[var] = {frame.id for frame in self.binder.frames} if isinstance(lvalue, MemberExpr) and self.inferred_attribute_types is not None: # Store inferred attribute type so that we can check consistency afterwards. if lvalue.def_var is not None: diff --git a/test-data/unit/check-optional.test b/test-data/unit/check-optional.test index 0ff7c2bc5b724..c6d3d1f486499 100644 --- a/test-data/unit/check-optional.test +++ b/test-data/unit/check-optional.test @@ -996,6 +996,29 @@ def f20(b: bool) -> None: x = None reveal_type(x) # N: Revealed type is "Any" +def f_unannot(): pass + +def f21(b: bool) -> None: + if b: + x = f_unannot() + else: + x = None + reveal_type(x) # N: Revealed type is "Any" + +def f22(b: bool) -> None: + if b: + z = 10 + if not b: + z = None # E: Incompatible types in assignment (expression has type "None", variable has type "int") + +def f23(b: bool) -> None: + if b: + z = 10 + if b: + z = 11 + else: + z = None # E: Incompatible types in assignment (expression has type "None", variable has type "int") + [builtins fixtures/exception.pyi] [case testOptionalBackwards2]