[use before def] handle class and function definitions

mypy/partially_defined.py

@@ -7,6 +7,7 @@
     AssignmentExpr,
     AssignmentStmt,
     BreakStmt,
+    ClassDef,
     Context,
     ContinueStmt,
     DictionaryComprehension,
@@ -258,13 +259,16 @@ def variable_may_be_undefined(self, name: str, context: Context) -> None:
         if self.msg.errors.is_error_code_enabled(errorcodes.PARTIALLY_DEFINED):
             self.msg.variable_may_be_undefined(name, context)
 
+    def process_definition(self, name: str) -> None:
+        # Was this name previously used? If yes, it's a use-before-definition error.
+        refs = self.tracker.pop_undefined_ref(name)
+        for ref in refs:
+            self.var_used_before_def(name, ref)
+        self.tracker.record_definition(name)
+
     def process_lvalue(self, lvalue: Lvalue | None) -> None:
         if isinstance(lvalue, NameExpr):
-            # Was this name previously used? If yes, it's a use-before-definition error.
-            refs = self.tracker.pop_undefined_ref(lvalue.name)
-            for ref in refs:
-                self.var_used_before_def(lvalue.name, ref)
-            self.tracker.record_definition(lvalue.name)
+            self.process_definition(lvalue.name)
         elif isinstance(lvalue, StarExpr):
             self.process_lvalue(lvalue.expr)
         elif isinstance(lvalue, (ListExpr, TupleExpr)):
@@ -314,7 +318,7 @@ def visit_match_stmt(self, o: MatchStmt) -> None:
         self.tracker.end_branch_statement()
 
     def visit_func_def(self, o: FuncDef) -> None:
-        self.tracker.record_definition(o.name)
+        self.process_definition(o.name)
         self.tracker.enter_scope()
         super().visit_func_def(o)
         self.tracker.exit_scope()
@@ -435,6 +439,10 @@ def visit_with_stmt(self, o: WithStmt) -> None:
             self.process_lvalue(idx)
         o.body.accept(self)
 
+    def visit_class_def(self, o: ClassDef) -> None:
+        self.process_definition(o.name)
+        super().visit_class_def(o)
+
     def visit_import(self, o: Import) -> None:
         for mod, alias in o.ids:
             if alias is not None:

test-data/unit/check-partially-defined.test

@@ -140,6 +140,17 @@ def f0(b: bool) -> None:
         fn = lambda: 2
     y = fn   # E: Name "fn" may be undefined
 
+[case testUseBeforeDefClass]
+# flags: --enable-error-code partially-defined --enable-error-code use-before-def
+def f(x: A):  # No error here.
+    pass
+y = A()  # E: Name "A" is used before definition
+class A: pass
+
+[case testUseBeforeDefFunc]
+# flags: --enable-error-code partially-defined --enable-error-code use-before-def
+foo() # E: Name "foo" is used before definition
+def foo(): pass
 [case testGenerator]
 # flags: --enable-error-code partially-defined
 if int():