[pyflakes] Allow forward references in class bases in stub files (F821)

crates/ruff_linter/resources/test/fixtures/pyflakes/F821_26.pyi

@@ -1,6 +1,6 @@
 """Tests for constructs allowed in `.pyi` stub files but not at runtime"""
 
-from typing import Optional, TypeAlias, Union
+from typing import Generic, NewType, Optional, TypeAlias, TypeVar, Union
 
 __version__: str
 __author__: str
@@ -33,6 +33,19 @@ class Leaf: ...
 class Tree(list[Tree | Leaf]): ...  # valid in a `.pyi` stub file, not in a `.py` runtime file
 class Tree2(list["Tree | Leaf"]): ...  # always okay
 
+# Generic bases can have forward references in stubs
+class Foo(Generic[T]): ...
+T = TypeVar("T")
+class Bar(Foo[Baz]): ...
+class Baz: ...
+
+# bases in general can be forward references in stubs
+class Eggs(Spam): ...
+class Spam: ...
+
+# NewType can have forward references
+MyNew = NewType("MyNew", MyClass)
+
 # Annotations are treated as assignments in .pyi files, but not in .py files
 class MyClass:
     foo: int
@@ -42,3 +55,6 @@ class MyClass:
 baz: MyClass
 eggs = baz  # valid in a `.pyi` stub file, not in a `.py` runtime file
 eggs = "baz"  # always okay
+
+class Blah:
+    class Blah2(Blah): ...

crates/ruff_linter/src/checkers/ast/deferred.rs

@@ -12,6 +12,8 @@ pub(crate) struct Visit<'a> {
     pub(crate) type_param_definitions: Vec<(&'a Expr, Snapshot)>,
     pub(crate) functions: Vec<Snapshot>,
     pub(crate) lambdas: Vec<Snapshot>,
+    /// N.B. This field should always be empty unless it's a stub file
+    pub(crate) class_bases: Vec<(&'a Expr, Snapshot)>,
 }
 
 impl Visit<'_> {
@@ -22,6 +24,7 @@ impl Visit<'_> {
             && self.type_param_definitions.is_empty()
             && self.functions.is_empty()
             && self.lambdas.is_empty()
+            && self.class_bases.is_empty()
     }
 }
 

crates/ruff_linter/src/checkers/ast/mod.rs

@@ -711,7 +711,9 @@ impl<'a> Visitor<'a> for Checker<'a> {
                 }
 
                 if let Some(arguments) = arguments {
+                    self.semantic.flags |= SemanticModelFlags::CLASS_BASE;
                     self.visit_arguments(arguments);
+                    self.semantic.flags -= SemanticModelFlags::CLASS_BASE;
                 }
 
                 let definition = docstrings::extraction::extract_definition(
@@ -934,6 +936,16 @@ impl<'a> Visitor<'a> for Checker<'a> {
 
     fn visit_expr(&mut self, expr: &'a Expr) {
         // Step 0: Pre-processing
+        if self.source_type.is_stub()
+            && self.semantic.in_class_base()
+            && !self.semantic.in_deferred_class_base()
+        {
+            self.visit
+                .class_bases
+                .push((expr, self.semantic.snapshot()));
+            return;
+        }
+
         if !self.semantic.in_typing_literal()
             // `in_deferred_type_definition()` will only be `true` if we're now visiting the deferred nodes
             // after having already traversed the source tree once. If we're now visiting the deferred nodes,
@@ -1967,6 +1979,33 @@ impl<'a> Checker<'a> {
         scope.add(id, binding_id);
     }
 
+    /// After initial traversal of the AST, visit all class bases that were deferred.
+    ///
+    /// This method should only be relevant in stub files, where forward references are
+    /// legal in class bases. For other kinds of Python files, using a forward reference
+    /// in a class base is never legal, so `self.visit.class_bases` should always be empty.
+    ///
+    /// For example, in a stub file:
+    /// ```python
+    /// class Foo(list[Bar]): ...  # <-- `Bar` is a forward reference in a class base
+    /// class Bar: ...
+    /// ```
+    fn visit_deferred_class_bases(&mut self) {
+        let snapshot = self.semantic.snapshot();
+        let deferred_bases = std::mem::take(&mut self.visit.class_bases);
+        debug_assert!(
+            self.source_type.is_stub() || deferred_bases.is_empty(),
+            "Class bases should never be deferred outside of stub files"
+        );
+        for (expr, snapshot) in deferred_bases {
+            self.semantic.restore(snapshot);
+            // Set this flag to avoid infinite recursion, or we'll just defer it again:
+            self.semantic.flags |= SemanticModelFlags::DEFERRED_CLASS_BASE;
+            self.visit_expr(expr);
+        }
+        self.semantic.restore(snapshot);
+    }
+
     /// After initial traversal of the AST, visit all "future type definitions".
     ///
     /// A "future type definition" is a type definition where [PEP 563] semantics
@@ -2157,6 +2196,7 @@ impl<'a> Checker<'a> {
     /// This includes lambdas, functions, type parameters, and type annotations.
     fn visit_deferred(&mut self, allocator: &'a typed_arena::Arena<Expr>) {
         while !self.visit.is_empty() {
+            self.visit_deferred_class_bases();
             self.visit_deferred_functions();
             self.visit_deferred_type_param_definitions();
             self.visit_deferred_lambdas();

crates/ruff_python_semantic/src/model.rs

@@ -1608,6 +1608,20 @@ impl<'a> SemanticModel<'a> {
             .intersects(SemanticModelFlags::DUNDER_ALL_DEFINITION)
     }
 
+    /// Return `true` if the model is visiting an item in a class's bases tuple
+    /// (e.g. `Foo` in `class Bar(Foo): ...`)
+    pub const fn in_class_base(&self) -> bool {
+        self.flags.intersects(SemanticModelFlags::CLASS_BASE)
+    }
+
+    /// Return `true` if the model is visiting an item in a class's bases tuple
+    /// that was initially deferred while traversing the AST.
+    /// (This only happens in stub files.)
+    pub const fn in_deferred_class_base(&self) -> bool {
+        self.flags
+            .intersects(SemanticModelFlags::DEFERRED_CLASS_BASE)
+    }
+
     /// Return an iterator over all bindings shadowed by the given [`BindingId`], within the
     /// containing scope, and across scopes.
     pub fn shadowed_bindings(
@@ -2021,6 +2035,20 @@ bitflags! {
         /// ```
         const F_STRING_REPLACEMENT_FIELD = 1 << 23;
 
+        /// The model is visiting the bases tuple of a class.
+        ///
+        /// For example, the model could be visiting `Foo` or `Bar` in:
+        ///
+        /// ```python
+        /// class Baz(Foo, Bar):
+        ///     pass
+        /// ```
+        const CLASS_BASE = 1 << 24;
+
+        /// The model is visiting a class base that was initially deferred
+        /// while traversing the AST. (This only happens in stub files.)
+        const DEFERRED_CLASS_BASE = 1 << 25;
+
         /// The context is in any type annotation.
         const ANNOTATION = Self::TYPING_ONLY_ANNOTATION.bits() | Self::RUNTIME_EVALUATED_ANNOTATION.bits() | Self::RUNTIME_REQUIRED_ANNOTATION.bits();