diff --git a/mypy/binder.py b/mypy/binder.py index 1c711ce9c6318..3c0c4e793359a 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, conditional_frame: bool = False) -> None: + self.id = id self.types: Dict[Key, Type] = {} self.unreachable = 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 @@ -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 @@ -109,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() + f = Frame(self._get_id(), conditional_frame) 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, + conditional_frame: bool = False, try_frame: bool = False) -> Iterator[Frame]: """Return a context manager that pushes/pops frames on enter/exit. @@ -401,7 +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() + 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 ba020f5d97d5b..a6b732dc025ec 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,14 @@ class TypeChecker(NodeVisitor[None], CheckerPluginInterface): # directly or indirectly. module_refs: Set[str] + # 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. plugin: Plugin @@ -229,6 +237,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() @@ -411,7 +420,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): @@ -2167,6 +2176,31 @@ 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 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. + has_if_ancestor = False + for frame in reversed(self.binder.frames): + if frame.id in decl_frame_map: + has_if_ancestor = frame.conditional_frame + 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 +3026,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: + # 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: @@ -3298,7 +3335,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, conditional_frame=True, fall_through=0): for e, b in zip(s.expr, s.body): t = get_proper_type(self.expr_checker.accept(e)) @@ -3437,7 +3474,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) @@ -4920,6 +4957,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 eedd5899c9497..c6d3d1f486499 100644 --- a/test-data/unit/check-optional.test +++ b/test-data/unit/check-optional.test @@ -832,3 +832,248 @@ 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 testOptionalBackwards1] +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: + 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 + +def f12(b: bool, a: Any) -> None: + if b: + z = a + 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 + +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]" + +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" + +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] + +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