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

## Summary

Fixes #3011.

Type checkers currently allow forward references in all contexts in stub
files, and stubs frequently make use of this capability (although it
doesn't actually seem to be specc'd anywhere --neither in PEP 484, nor
https://typing.readthedocs.io/en/latest/source/stubs.html#id6, nor the
CPython typing docs). Implementing it so that Ruff allows forward
references in _all contexts_ in stub files seems non-trivial, however
(or at least, I couldn't figure out how to do it easily), so this PR
does not do that. Perhaps it _should_; if we think this apporach isn't
principled enough, I'm happy to close it and postpone changing anything
here.

However, this does reduce the number of F821 errors Ruff emits on
typeshed down from 76 to 2, which would mean that we could enable the
rule at typeshed. The remaining 2 F821 errors can be trivially fixed at
typeshed by moving definitions around; forward references in class bases
were really the only remaining places where there was a real _use case_
for forward references in stub files that Ruff wasn't yet allowing.

## Test plan

`cargo test`. I also ran this PR branch on typeshed to check to see if
there were any new false positives caused by the changes here; there
were none.

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();