Support making a variable Optional in an else branch (#11002)

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.

Author: msullivan

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)

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)
@@ -4925,6 +4962,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)