preserve backwards compatibility

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

@@ -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 Optional, TypeAlias, Union, TYPE_CHECKING
 
 __version__: str
 __author__: str
@@ -42,3 +42,7 @@ class MyClass:
 baz: MyClass
 eggs = baz  # valid in a `.pyi` stub file, not in a `.py` runtime file
 eggs = "baz"  # always okay
+
+if TYPE_CHECKING:
+    class Tree3(list[Tree3 | Leaf]): ...  # Always okay if it's in a `TYPE_CHECKING` block
+    Recursive: TypeAlias = Recursive | None  # Always okay if it's in a `TYPE_CHECKING` block

crates/ruff_linter/resources/test/fixtures/pyflakes/F821_27.py

@@ -1,7 +1,7 @@
-"""Tests for constructs allowed when `__future__` annotations are enabled but not otherwise"""
+"""Tests for forward references (some allowed, some still not) in the context of `from __future__ import annotations`"""
 from __future__ import annotations
 
-from typing import Optional, TypeAlias, Union
+from typing import Optional, TypeAlias, Union, TYPE_CHECKING
 
 __version__: str
 __author__: str
@@ -45,4 +45,8 @@ class D: ...
 # More circular references
 class Leaf: ...
 class Tree(list[Tree | Leaf]): ...  # Still invalid even when `__future__.annotations` are enabled
-class Tree2(list["Tree | Leaf"]): ...  # always okay
+class Tree2(list["Tree2 | Leaf"]): ...  # always okay
+
+if TYPE_CHECKING:
+    class Tree3(list[Tree3 | Leaf]): ...  # Always okay if it's in a `TYPE_CHECKING` block
+    Recursive: TypeAlias = Recursive | None  # Always okay if it's in a `TYPE_CHECKING` block

crates/ruff_linter/src/checkers/ast/analyze/expression.rs

@@ -32,7 +32,8 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
                 if let Some(operator) = typing::to_pep604_operator(value, slice, &checker.semantic)
                 {
                     if checker.enabled(Rule::FutureRewritableTypeAnnotation) {
-                        if !checker.semantic.future_annotations_or_stub()
+                        if !checker.semantic.in_typing_only_context()
+                            && !checker.semantic.future_annotations()
                             && checker.settings.target_version < PythonVersion::Py310
                             && checker.settings.target_version >= PythonVersion::Py37
                             && checker.semantic.in_annotation()
@@ -44,10 +45,10 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
                         }
                     }
                     if checker.enabled(Rule::NonPEP604Annotation) {
-                        if checker.source_type.is_stub()
+                        if checker.semantic.in_typing_only_context()
                             || checker.settings.target_version >= PythonVersion::Py310
                             || (checker.settings.target_version >= PythonVersion::Py37
-                                && checker.semantic.future_annotations_or_stub()
+                                && checker.semantic.future_annotations()
                                 && checker.semantic.in_annotation()
                                 && !checker.settings.pyupgrade.keep_runtime_typing)
                         {
@@ -59,7 +60,8 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
 
             // Ex) list[...]
             if checker.enabled(Rule::FutureRequiredTypeAnnotation) {
-                if !checker.semantic.future_annotations_or_stub()
+                if !checker.semantic.in_typing_only_context()
+                    && !checker.semantic.future_annotations()
                     && checker.settings.target_version < PythonVersion::Py39
                     && checker.semantic.in_annotation()
                     && typing::is_pep585_generic(value, &checker.semantic)
@@ -189,7 +191,8 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
                             typing::to_pep585_generic(expr, &checker.semantic)
                         {
                             if checker.enabled(Rule::FutureRewritableTypeAnnotation) {
-                                if !checker.semantic.future_annotations_or_stub()
+                                if !checker.semantic.in_typing_only_context()
+                                    && !checker.semantic.future_annotations()
                                     && checker.settings.target_version < PythonVersion::Py39
                                     && checker.settings.target_version >= PythonVersion::Py37
                                     && checker.semantic.in_annotation()
@@ -199,10 +202,10 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
                                 }
                             }
                             if checker.enabled(Rule::NonPEP585Annotation) {
-                                if checker.source_type.is_stub()
+                                if checker.semantic.in_typing_only_context()
                                     || checker.settings.target_version >= PythonVersion::Py39
                                     || (checker.settings.target_version >= PythonVersion::Py37
-                                        && checker.semantic.future_annotations_or_stub()
+                                        && checker.semantic.future_annotations()
                                         && checker.semantic.in_annotation()
                                         && !checker.settings.pyupgrade.keep_runtime_typing)
                                 {
@@ -272,7 +275,8 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
             ]) {
                 if let Some(replacement) = typing::to_pep585_generic(expr, &checker.semantic) {
                     if checker.enabled(Rule::FutureRewritableTypeAnnotation) {
-                        if !checker.semantic.future_annotations_or_stub()
+                        if !checker.semantic.in_typing_only_context()
+                            && !checker.semantic.future_annotations()
                             && checker.settings.target_version < PythonVersion::Py39
                             && checker.settings.target_version >= PythonVersion::Py37
                             && checker.semantic.in_annotation()
@@ -284,10 +288,10 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
                         }
                     }
                     if checker.enabled(Rule::NonPEP585Annotation) {
-                        if checker.source_type.is_stub()
+                        if checker.semantic.in_typing_only_context()
                             || checker.settings.target_version >= PythonVersion::Py39
                             || (checker.settings.target_version >= PythonVersion::Py37
-                                && checker.semantic.future_annotations_or_stub()
+                                && checker.semantic.future_annotations()
                                 && checker.semantic.in_annotation()
                                 && !checker.settings.pyupgrade.keep_runtime_typing)
                         {
@@ -1177,7 +1181,8 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
         }) => {
             // Ex) `str | None`
             if checker.enabled(Rule::FutureRequiredTypeAnnotation) {
-                if !checker.semantic.future_annotations_or_stub()
+                if !checker.semantic.in_typing_only_context()
+                    && !checker.semantic.future_annotations()
                     && checker.settings.target_version < PythonVersion::Py310
                     && checker.semantic.in_annotation()
                 {

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

@@ -56,10 +56,11 @@ impl AnnotationContext {
             _ => {}
         }
 
-        // If `__future__` annotations are enabled or it's a stub file,
+        // If `__future__` annotations are enabled or it's a stub file
+        // or we're in an `if TYPE_CHECKING` block,
         // then annotations are never evaluated at runtime,
         // so we can treat them as typing-only.
-        if semantic.future_annotations_or_stub() {
+        if semantic.in_typing_only_context() || semantic.future_annotations() {
             return Self::TypingOnly;
         }
 
@@ -88,7 +89,7 @@ impl AnnotationContext {
             semantic,
         ) {
             Self::RuntimeRequired
-        } else if semantic.future_annotations_or_stub() {
+        } else if semantic.in_typing_only_context() || semantic.future_annotations() {
             Self::TypingOnly
         } else {
             Self::RuntimeEvaluated

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

@@ -893,7 +893,7 @@ 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()
+        if self.semantic.in_typing_only_context()
             && self.semantic.in_class_base()
             && !self.semantic.in_deferred_class_base()
         {
@@ -903,15 +903,16 @@ impl<'a> Visitor<'a> for Checker<'a> {
             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,
-            // we can't defer again, or we'll infinitely recurse!
+        // `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,
+        // we can't defer again, or we'll infinitely recurse!
+        let defer_probable_type_definition = self.semantic.in_type_definition()
             && !self.semantic.in_deferred_type_definition()
-            && self.semantic.in_type_definition()
-            && self.semantic.future_annotations_or_stub()
-            && (self.semantic.in_annotation() || self.source_type.is_stub())
-        {
+            && !self.semantic.in_typing_literal()
+            && (self.semantic.in_typing_only_context()
+                || (self.semantic.future_annotations() && self.semantic.in_annotation()));
+
+        if defer_probable_type_definition {
             if let Expr::StringLiteral(ast::ExprStringLiteral { value, .. }) = expr {
                 self.visit.string_type_definitions.push((
                     expr.range(),
@@ -1953,9 +1954,10 @@ impl<'a> Checker<'a> {
 
     /// 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.
+    /// This method should only be relevant in stub files or `TYPE_CHECKING` blocks,
+    /// where forward references are legal in class bases. In other contexts,
+    /// 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
@@ -1965,12 +1967,9 @@ impl<'a> Checker<'a> {
     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);
+            debug_assert!(self.semantic.in_typing_only_context());
             // 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);
@@ -1981,8 +1980,9 @@ impl<'a> Checker<'a> {
     /// After initial traversal of the AST, visit all "future type definitions".
     ///
     /// A "future type definition" is a type definition where [PEP 563] semantics
-    /// apply (i.e., an annotation in a module that has `from __future__ import annotations`
-    /// at the top of the file, or an annotation in a stub file). These type definitions
+    /// apply (i.e., a definition in a module that has `from __future__ import annotations`
+    /// at the top of the file, a definition in a `TYPE_CHECKING` block,
+    /// or an annotation in a stub file). These type definitions
     /// support forward references, so they are deferred on initial traversal
     /// of the source tree.
     ///
@@ -2003,15 +2003,6 @@ impl<'a> Checker<'a> {
             let type_definitions = std::mem::take(&mut self.visit.future_type_definitions);
             for (expr, snapshot) in type_definitions {
                 self.semantic.restore(snapshot);
-
-                // Type definitions should only be considered "`__future__` type definitions"
-                // if they are annotations in a module where `from __future__ import
-                // annotations` is active, or they are type definitions in a stub file.
-                debug_assert!(
-                    self.semantic.future_annotations_or_stub()
-                        && (self.source_type.is_stub() || self.semantic.in_annotation())
-                );
-
                 self.semantic.flags |= SemanticModelFlags::TYPE_DEFINITION
                     | SemanticModelFlags::FUTURE_TYPE_DEFINITION;
                 self.visit_expr(expr);
@@ -2071,7 +2062,10 @@ impl<'a> Checker<'a> {
 
                     self.semantic.restore(snapshot);
 
-                    if self.semantic.in_annotation() && self.semantic.future_annotations_or_stub() {
+                    if self.semantic.in_annotation()
+                        && (self.semantic.in_typing_only_context()
+                            || self.semantic.future_annotations())
+                    {
                         if self.enabled(Rule::QuotedAnnotation) {
                             pyupgrade::rules::quoted_annotation(self, value, range);
                         }

crates/ruff_linter/src/rules/pyflakes/snapshots/ruff_linter__rules__pyflakes__tests__F821_F821_27.py.snap

@@ -42,5 +42,5 @@ F821_27.py:47:17: F821 Undefined name `Tree`
 46 | class Leaf: ...
 47 | class Tree(list[Tree | Leaf]): ...  # Still invalid even when `__future__.annotations` are enabled
    |                 ^^^^ F821
-48 | class Tree2(list["Tree | Leaf"]): ...  # always okay
+48 | class Tree2(list["Tree2 | Leaf"]): ...  # always okay
    |

crates/ruff_python_semantic/src/model.rs

@@ -1653,10 +1653,30 @@ impl<'a> SemanticModel<'a> {
             .intersects(SemanticModelFlags::MODULE_DOCSTRING_BOUNDARY)
     }
 
-    /// Return `true` if `__future__`-style type annotations are enabled.
-    pub const fn future_annotations_or_stub(&self) -> bool {
+    /// Return `true` if we're in a file with `from __future__ import annotations`
+    /// at the top of the file.
+    ///
+    /// N.B. Checking this flag is *not* generally sufficient to determine whether
+    /// `__future__`-style forward references are enabled, since forward references
+    /// are allowed in two other contexts even without the
+    /// `from __future__ import annotations` import: stub files and `if TYPE_CHECKING`
+    /// blocks. To determine whether forward references are legal in a given context,
+    /// you'll need to do something like:
+    ///
+    /// ```ignore
+    /// if semantic.in_typing_only_context() || (
+    ///     semantic.future_annotations() && semantic.in_annotation()
+    /// ) {}
+    /// ```
+    pub const fn future_annotations(&self) -> bool {
         self.flags
-            .intersects(SemanticModelFlags::FUTURE_ANNOTATIONS_OR_STUB)
+            .intersects(SemanticModelFlags::FUTURE_ANNOTATIONS)
+    }
+
+    /// Return `true` if we're in a block of code that is never executed at runtime:
+    /// either a stub file, or a `TYPE_CHECKING` block
+    pub const fn in_typing_only_context(&self) -> bool {
+        self.flags.intersects(SemanticModelFlags::TYPING_ONLY_BLOCK)
     }
 
     /// Return `true` if the model is in a stub file (i.e., a file with a `.pyi` extension).
@@ -2015,11 +2035,9 @@ bitflags! {
         /// The model is in a Python stub file (i.e., a `.pyi` file).
         const STUB_FILE = 1 << 16;
 
-        /// `__future__`-style type annotations are enabled in this model.
-        /// That could be because it's a stub file,
-        /// or it could be because it's a non-stub file that has `from __future__ import annotations`
-        /// a the top of the module.
-        const FUTURE_ANNOTATIONS_OR_STUB = Self::FUTURE_ANNOTATIONS.bits() | Self::STUB_FILE.bits();
+        /// The model is in a context that's never evaluated at runtime:
+        /// either a stub file or a type-checking block
+        const TYPING_ONLY_BLOCK = Self::STUB_FILE.bits() | Self::TYPE_CHECKING_BLOCK.bits();
 
         /// The model has traversed past the module docstring.
         ///