Implement foundation for detecting partially defined vars

Author: ilinum

mypy/build.py

@@ -47,9 +47,11 @@
 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.semanal import SemanticAnalyzer
 from mypy.semanal_pass1 import SemanticAnalyzerPreAnalysis
+from mypy.undefined_vars import UndefinedVariableVisitor
 from mypy.util import (
     DecodeError,
     decode_python_encoding,
@@ -2340,6 +2342,11 @@ def finish_passes(self) -> None:
         manager = self.manager
         if self.options.semantic_analysis_only:
             return
+        if manager.options.disallow_undefined_vars:
+            manager.errors.set_file(self.xpath, self.tree.fullname, options=manager.options)
+            self.tree.accept(
+                UndefinedVariableVisitor(MessageBuilder(manager.errors, manager.modules))
+            )
         t0 = time_ref()
         with self.wrap_context():
             # Some tests (and tools) want to look at the set of all types.

mypy/main.py

@@ -823,6 +823,8 @@ def add_invertible_flag(
         help="Make arguments prepended via Concatenate be truly positional-only",
         group=strictness_group,
     )
+    # Experiment flag to detect undefined variables being used.
+    add_invertible_flag("--disallow-undefined-vars", default=False, help=argparse.SUPPRESS)
 
     strict_help = "Strict mode; enables the following flags: {}".format(
         ", ".join(strict_flag_names)

mypy/messages.py

@@ -1216,6 +1216,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)
+
     def first_argument_for_super_must_be_type(self, actual: Type, context: Context) -> None:
         actual = get_proper_type(actual)
         if isinstance(actual, Instance):

mypy/options.py

@@ -180,6 +180,9 @@ def __init__(self) -> None:
         # Make arguments prepended via Concatenate be truly positional-only.
         self.strict_concatenate = False
 
+        # Disallow using vars that could be undefined.
+        self.disallow_undefined_vars = False
+
         # Report an error for any branches inferred to be unreachable as a result of
         # type analysis.
         self.warn_unreachable = False

mypy/undefined_vars.py

@@ -0,0 +1,183 @@
+from __future__ import annotations
+
+from typing import NamedTuple
+
+from mypy.messages import MessageBuilder
+from mypy.nodes import (
+    AssignmentStmt,
+    FuncDef,
+    FuncItem,
+    IfStmt,
+    ListExpr,
+    Lvalue,
+    NameExpr,
+    TupleExpr,
+    WhileStmt,
+)
+from mypy.traverser import TraverserVisitor
+
+
+class DefinedVariables(NamedTuple):
+    """DefinedVariables 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 BranchingTracker:
+    def __init__(self) -> None:
+        self.defined_by_branch: list[DefinedVariables] = [
+            DefinedVariables(may_be_defined=set(), must_be_defined=set())
+        ]
+
+    def next_branch(self) -> None:
+        self.defined_by_branch.append(
+            DefinedVariables(may_be_defined=set(), must_be_defined=set())
+        )
+
+    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: DefinedVariables) -> 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) -> DefinedVariables:
+        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 DefinedVariables(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:
+        # todo(stas): we should initialize this with some variables.
+        # There's always at least one scope. Within each scope, there's at least one "global" BranchingTracker.
+        self.scopes: list[list[BranchingTracker]] = [[BranchingTracker()]]
+
+    def _scope(self) -> list[BranchingTracker]:
+        assert len(self.scopes) > 0
+        return self.scopes[-1]
+
+    def enter_scope(self) -> None:
+        self.scopes.append([BranchingTracker()])
+
+    def exit_scope(self) -> None:
+        self.scopes.pop()
+
+    def start_branch_statement(self) -> None:
+        self._scope().append(BranchingTracker())
+
+    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 UndefinedVariableVisitor(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_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)

test-data/unit/check-undefined-vars.test

@@ -0,0 +1,121 @@
+[case testDefinedInOneBranch]
+# flags: --disallow-undefined-vars
+if int():
+    a = 1
+else:
+    x = 2
+z = a + 1  # E: Name "a" may be undefined
+[case testElif]
+# flags: --disallow-undefined-vars
+if int():
+    a = 1
+elif int():
+    a = 2
+else:
+    x = 3
+
+z = a + 1  # E: Name "a" may be undefined
+
+[case testDefinedInAllBranches]
+# flags: --disallow-undefined-vars
+if int():
+    a = 1
+elif int():
+    a = 2
+else:
+    a = 3
+z = a + 1
+
+[case testOmittedElse]
+# flags: --disallow-undefined-vars
+if int():
+    a = 1
+z = a + 1  # E: Name "a" may be undefined
+
+[case testUpdatedInIf]
+# flags: --disallow-undefined-vars
+# Variable a is already defined. Just updating it in an "if" is acceptable.
+a = 1
+if int():
+    a = 2
+z = a + 1
+
+[case testNestedIf]
+# flags: --disallow-undefined-vars
+if int():
+    if int():
+        a = 1
+        x = 1
+        x = x + 1
+    else:
+        a = 2
+    b = a + x  # E: Name "x" may be undefined
+    b = b + 1
+else:
+    b = 2
+z = a + b  # E: Name "a" may be undefined
+
+[case testVeryNestedIf]
+# flags: --disallow-undefined-vars
+if int():
+    if int():
+        if int():
+            a = 1
+        else:
+            a = 2
+        x = a
+    else:
+        a = 2
+    b = a
+else:
+    b = 2
+z = a + b  # E: Name "a" may be undefined
+
+[case testTupleUnpack]
+# flags: --disallow-undefined-vars
+
+if int():
+    (x, y) = (1, 2)
+else:
+    [y, z] = [1, 2]
+a = y + x  # E: Name "x" may be undefined
+a = y + z  # E: Name "z" may be undefined
+
+[case testRedefined]
+# flags: --disallow-undefined-vars
+
+if int():
+    y = 2
+y = 3
+x = y + 2
+
+[case testScope]
+# flags: --disallow-undefined-vars
+def foo() -> None:
+    if int():
+        y = 2
+
+if int():
+    y = 3
+x = y  # E: Name "y" may be undefined
+
+[case testFuncParams]
+# flags: --disallow-undefined-vars
+def foo(a: int) -> None:
+    if int():
+        a = 2
+    x = a
+
+[case testWhile]
+# flags: --disallow-undefined-vars
+while int():
+    x = 1
+
+y = x  # E: Name "x" may be undefined
+
+while int():
+    z = 1
+else:
+    z = 2
+
+y = z  # No error.