Implement foundation for detecting partially defined vars (#13601)

This diff lays the foundation for detecting partially defined variables.
Think of the following situation:
```
if foo():
    x = 1
print(x)  # Error: "x" may be undefined.
```

Now, mypy will generate the error in such a case.

Note that this diff is not complete. It still generates a lot of false
positives. Running it on mypy itself generated 182 errors.
Therefore, this feature is disabled by default and the error code
must be explicitly enabled. I will implement it in multiple PRs.

Author: ilinum

mypy/build.py

@@ -47,7 +47,9 @@
 from mypy.checker import TypeChecker
 from mypy.errors import CompileError, ErrorInfo, Errors, report_internal_error
 from mypy.indirection import TypeIndirectionVisitor
+from mypy.messages import MessageBuilder
 from mypy.nodes import Import, ImportAll, ImportBase, ImportFrom, MypyFile, SymbolTable
+from mypy.partially_defined import PartiallyDefinedVariableVisitor
 from mypy.semanal import SemanticAnalyzer
 from mypy.semanal_pass1 import SemanticAnalyzerPreAnalysis
 from mypy.util import (
@@ -2335,6 +2337,15 @@ def type_check_second_pass(self) -> bool:
         self.time_spent_us += time_spent_us(t0)
         return result
 
+    def detect_partially_defined_vars(self) -> None:
+        assert self.tree is not None, "Internal error: method must be called on parsed file only"
+        manager = self.manager
+        if manager.errors.is_error_code_enabled(codes.PARTIALLY_DEFINED):
+            manager.errors.set_file(self.xpath, self.tree.fullname, options=manager.options)
+            self.tree.accept(
+                PartiallyDefinedVariableVisitor(MessageBuilder(manager.errors, manager.modules))
+            )
+
     def finish_passes(self) -> None:
         assert self.tree is not None, "Internal error: method must be called on parsed file only"
         manager = self.manager
@@ -3364,6 +3375,7 @@ def process_stale_scc(graph: Graph, scc: list[str], manager: BuildManager) -> No
         graph[id].type_check_first_pass()
         if not graph[id].type_checker().deferred_nodes:
             unfinished_modules.discard(id)
+            graph[id].detect_partially_defined_vars()
             graph[id].finish_passes()
 
     while unfinished_modules:
@@ -3372,6 +3384,7 @@ def process_stale_scc(graph: Graph, scc: list[str], manager: BuildManager) -> No
                 continue
             if not graph[id].type_check_second_pass():
                 unfinished_modules.discard(id)
+                graph[id].detect_partially_defined_vars()
                 graph[id].finish_passes()
     for id in stale:
         graph[id].generate_unused_ignore_notes()

mypy/errorcodes.py

@@ -124,6 +124,12 @@ def __str__(self) -> str:
 UNREACHABLE: Final = ErrorCode(
     "unreachable", "Warn about unreachable statements or expressions", "General"
 )
+PARTIALLY_DEFINED: Final[ErrorCode] = ErrorCode(
+    "partially-defined",
+    "Warn about variables that are defined only in some execution paths",
+    "General",
+    default_enabled=False,
+)
 REDUNDANT_EXPR: Final = ErrorCode(
     "redundant-expr", "Warn about redundant expressions", "General", default_enabled=False
 )

mypy/messages.py

@@ -1214,6 +1214,9 @@ def invalid_keyword_var_arg(self, typ: Type, is_mapping: bool, context: Context)
     def undefined_in_superclass(self, member: str, context: Context) -> None:
         self.fail(f'"{member}" undefined in superclass', context)
 
+    def variable_may_be_undefined(self, name: str, context: Context) -> None:
+        self.fail(f'Name "{name}" may be undefined', context, code=codes.PARTIALLY_DEFINED)
+
     def first_argument_for_super_must_be_type(self, actual: Type, context: Context) -> None:
         actual = get_proper_type(actual)
         if isinstance(actual, Instance):

mypy/partially_defined.py

@@ -0,0 +1,201 @@
+from __future__ import annotations
+
+from typing import NamedTuple
+
+from mypy.messages import MessageBuilder
+from mypy.nodes import (
+    AssignmentStmt,
+    ForStmt,
+    FuncDef,
+    FuncItem,
+    IfStmt,
+    ListExpr,
+    Lvalue,
+    NameExpr,
+    TupleExpr,
+    WhileStmt,
+)
+from mypy.traverser import TraverserVisitor
+
+
+class DefinedVars(NamedTuple):
+    """DefinedVars contains information about variable definition at the end of a branching statement.
+    `if` and `match` are examples of branching statements.
+
+    `may_be_defined` contains variables that were defined in only some branches.
+    `must_be_defined` contains variables that were defined in all branches.
+    """
+
+    may_be_defined: set[str]
+    must_be_defined: set[str]
+
+
+class BranchStatement:
+    def __init__(self, already_defined: DefinedVars) -> None:
+        self.already_defined = already_defined
+        self.defined_by_branch: list[DefinedVars] = [
+            DefinedVars(may_be_defined=set(), must_be_defined=set(already_defined.must_be_defined))
+        ]
+
+    def next_branch(self) -> None:
+        self.defined_by_branch.append(
+            DefinedVars(
+                may_be_defined=set(), must_be_defined=set(self.already_defined.must_be_defined)
+            )
+        )
+
+    def record_definition(self, name: str) -> None:
+        assert len(self.defined_by_branch) > 0
+        self.defined_by_branch[-1].must_be_defined.add(name)
+        self.defined_by_branch[-1].may_be_defined.discard(name)
+
+    def record_nested_branch(self, vars: DefinedVars) -> None:
+        assert len(self.defined_by_branch) > 0
+        current_branch = self.defined_by_branch[-1]
+        current_branch.must_be_defined.update(vars.must_be_defined)
+        current_branch.may_be_defined.update(vars.may_be_defined)
+        current_branch.may_be_defined.difference_update(current_branch.must_be_defined)
+
+    def is_possibly_undefined(self, name: str) -> bool:
+        assert len(self.defined_by_branch) > 0
+        return name in self.defined_by_branch[-1].may_be_defined
+
+    def done(self) -> DefinedVars:
+        assert len(self.defined_by_branch) > 0
+        if len(self.defined_by_branch) == 1:
+            # If there's only one branch, then we just return current.
+            # Note that this case is a different case when an empty branch is omitted (e.g. `if` without `else`).
+            return self.defined_by_branch[0]
+
+        # must_be_defined is a union of must_be_defined of all branches.
+        must_be_defined = set(self.defined_by_branch[0].must_be_defined)
+        for branch_vars in self.defined_by_branch[1:]:
+            must_be_defined.intersection_update(branch_vars.must_be_defined)
+        # may_be_defined are all variables that are not must be defined.
+        all_vars = set()
+        for branch_vars in self.defined_by_branch:
+            all_vars.update(branch_vars.may_be_defined)
+            all_vars.update(branch_vars.must_be_defined)
+        may_be_defined = all_vars.difference(must_be_defined)
+        return DefinedVars(may_be_defined=may_be_defined, must_be_defined=must_be_defined)
+
+
+class DefinedVariableTracker:
+    """DefinedVariableTracker manages the state and scope for the UndefinedVariablesVisitor."""
+
+    def __init__(self) -> None:
+        # There's always at least one scope. Within each scope, there's at least one "global" BranchingStatement.
+        self.scopes: list[list[BranchStatement]] = [
+            [BranchStatement(DefinedVars(may_be_defined=set(), must_be_defined=set()))]
+        ]
+
+    def _scope(self) -> list[BranchStatement]:
+        assert len(self.scopes) > 0
+        return self.scopes[-1]
+
+    def enter_scope(self) -> None:
+        assert len(self._scope()) > 0
+        self.scopes.append([BranchStatement(self._scope()[-1].defined_by_branch[-1])])
+
+    def exit_scope(self) -> None:
+        self.scopes.pop()
+
+    def start_branch_statement(self) -> None:
+        assert len(self._scope()) > 0
+        self._scope().append(BranchStatement(self._scope()[-1].defined_by_branch[-1]))
+
+    def next_branch(self) -> None:
+        assert len(self._scope()) > 1
+        self._scope()[-1].next_branch()
+
+    def end_branch_statement(self) -> None:
+        assert len(self._scope()) > 1
+        result = self._scope().pop().done()
+        self._scope()[-1].record_nested_branch(result)
+
+    def record_declaration(self, name: str) -> None:
+        assert len(self.scopes) > 0
+        assert len(self.scopes[-1]) > 0
+        self._scope()[-1].record_definition(name)
+
+    def is_possibly_undefined(self, name: str) -> bool:
+        assert len(self._scope()) > 0
+        # A variable is undefined if it's in a set of `may_be_defined` but not in `must_be_defined`.
+        # Cases where a variable is not defined altogether are handled by semantic analyzer.
+        return self._scope()[-1].is_possibly_undefined(name)
+
+
+class PartiallyDefinedVariableVisitor(TraverserVisitor):
+    """Detect variables that are defined only part of the time.
+
+    This visitor detects the following case:
+    if foo():
+        x = 1
+    print(x)  # Error: "x" may be undefined.
+
+    Note that this code does not detect variables not defined in any of the branches -- that is
+    handled by the semantic analyzer.
+    """
+
+    def __init__(self, msg: MessageBuilder) -> None:
+        self.msg = msg
+        self.tracker = DefinedVariableTracker()
+
+    def process_lvalue(self, lvalue: Lvalue) -> None:
+        if isinstance(lvalue, NameExpr):
+            self.tracker.record_declaration(lvalue.name)
+        elif isinstance(lvalue, (ListExpr, TupleExpr)):
+            for item in lvalue.items:
+                self.process_lvalue(item)
+
+    def visit_assignment_stmt(self, o: AssignmentStmt) -> None:
+        for lvalue in o.lvalues:
+            self.process_lvalue(lvalue)
+        super().visit_assignment_stmt(o)
+
+    def visit_if_stmt(self, o: IfStmt) -> None:
+        for e in o.expr:
+            e.accept(self)
+        self.tracker.start_branch_statement()
+        for b in o.body:
+            b.accept(self)
+            self.tracker.next_branch()
+        if o.else_body:
+            o.else_body.accept(self)
+        self.tracker.end_branch_statement()
+
+    def visit_func_def(self, o: FuncDef) -> None:
+        self.tracker.enter_scope()
+        super().visit_func_def(o)
+        self.tracker.exit_scope()
+
+    def visit_func(self, o: FuncItem) -> None:
+        if o.arguments is not None:
+            for arg in o.arguments:
+                self.tracker.record_declaration(arg.variable.name)
+        super().visit_func(o)
+
+    def visit_for_stmt(self, o: ForStmt) -> None:
+        o.expr.accept(self)
+        self.process_lvalue(o.index)
+        o.index.accept(self)
+        self.tracker.start_branch_statement()
+        o.body.accept(self)
+        self.tracker.next_branch()
+        if o.else_body:
+            o.else_body.accept(self)
+        self.tracker.end_branch_statement()
+
+    def visit_while_stmt(self, o: WhileStmt) -> None:
+        o.expr.accept(self)
+        self.tracker.start_branch_statement()
+        o.body.accept(self)
+        self.tracker.next_branch()
+        if o.else_body:
+            o.else_body.accept(self)
+        self.tracker.end_branch_statement()
+
+    def visit_name_expr(self, o: NameExpr) -> None:
+        if self.tracker.is_possibly_undefined(o.name):
+            self.msg.variable_may_be_undefined(o.name, o)
+        super().visit_name_expr(o)

mypy/server/update.py

@@ -651,6 +651,7 @@ def restore(ids: list[str]) -> None:
     state.type_checker().reset()
     state.type_check_first_pass()
     state.type_check_second_pass()
+    state.detect_partially_defined_vars()
     t2 = time.time()
     state.finish_passes()
     t3 = time.time()