diff --git a/mypy/partially_defined.py b/mypy/partially_defined.py index 70a454beae9c..7e71af4044ed 100644 --- a/mypy/partially_defined.py +++ b/mypy/partially_defined.py @@ -234,6 +234,7 @@ class PartiallyDefinedVariableVisitor(ExtendedTraverserVisitor): def __init__(self, msg: MessageBuilder, type_map: dict[Expression, Type]) -> None: self.msg = msg self.type_map = type_map + self.loop_depth = 0 self.tracker = DefinedVariableTracker() def process_lvalue(self, lvalue: Lvalue | None) -> None: @@ -319,10 +320,12 @@ def visit_for_stmt(self, o: ForStmt) -> None: self.process_lvalue(o.index) o.index.accept(self) self.tracker.start_branch_statement() + self.loop_depth += 1 o.body.accept(self) self.tracker.next_branch() if o.else_body: o.else_body.accept(self) + self.loop_depth -= 1 self.tracker.end_branch_statement() def visit_return_stmt(self, o: ReturnStmt) -> None: @@ -354,7 +357,9 @@ def visit_expression_stmt(self, o: ExpressionStmt) -> None: def visit_while_stmt(self, o: WhileStmt) -> None: o.expr.accept(self) self.tracker.start_branch_statement() + self.loop_depth += 1 o.body.accept(self) + self.loop_depth -= 1 if not checker.is_true_literal(o.expr): self.tracker.next_branch() if o.else_body: @@ -380,7 +385,10 @@ def visit_name_expr(self, o: NameExpr) -> None: self.tracker.record_definition(o.name) elif self.tracker.is_defined_in_different_branch(o.name): # A variable is defined in one branch but used in a different branch. - self.msg.var_used_before_def(o.name, o) + if self.loop_depth > 0: + self.msg.variable_may_be_undefined(o.name, o) + else: + self.msg.var_used_before_def(o.name, o) elif self.tracker.is_undefined(o.name): # A variable is undefined. It could be due to two things: # 1. A variable is just totally undefined diff --git a/test-data/unit/check-partially-defined.test b/test-data/unit/check-partially-defined.test index c63023aa2746..2028362cedbe 100644 --- a/test-data/unit/check-partially-defined.test +++ b/test-data/unit/check-partially-defined.test @@ -206,6 +206,48 @@ def f5() -> int: return 3 return 1 +[case testDefinedDifferentBranchUseBeforeDef] +# flags: --enable-error-code partially-defined --enable-error-code use-before-def + +def f0() -> None: + if int(): + x = 0 + else: + y = x # E: Name "x" is used before definition + z = x # E: Name "x" is used before definition + +def f1() -> None: + x = 1 + if int(): + x = 0 + else: + y = x # No error. + + +[case testDefinedDifferentBranchPartiallyDefined] +# flags: --enable-error-code partially-defined --enable-error-code use-before-def + +def f0() -> None: + first_iter = True + for i in [0, 1]: + if first_iter: + first_iter = False + x = 0 + else: + # This is technically a false positive but mypy isn't smart enough for this yet. + y = x # E: Name "x" may be undefined + z = x # E: Name "x" may be undefined + + +def f1() -> None: + while True: + if int(): + x = 0 + else: + y = x # E: Name "x" may be undefined + z = x # E: Name "x" may be undefined + + [case testAssert] # flags: --enable-error-code partially-defined def f1() -> int: @@ -394,21 +436,7 @@ def f0() -> None: x = y # E: Name "y" is used before definition y: int = 1 -def f1() -> None: - if int(): - x = 0 - else: - y = x # E: Name "x" is used before definition - z = x # E: Name "x" is used before definition - def f2() -> None: - x = 1 - if int(): - x = 0 - else: - y = x # No error. - -def f3() -> None: if int(): pass else: @@ -418,14 +446,14 @@ def f3() -> None: def inner2() -> None: z = 0 -def f4() -> None: +def f3() -> None: if int(): pass else: y = z # E: Name "z" is used before definition z: int = 2 -def f5() -> None: +def f4() -> None: if int(): pass else: