Review feedback

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

@@ -0,0 +1,66 @@
+use ruff_python_semantic::{ScopeKind, SemanticModel};
+
+use crate::rules::flake8_type_checking;
+use crate::settings::LinterSettings;
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub(super) enum AnnotationContext {
+    /// Python will evaluate the annotation at runtime, but it's not _required_ and, as such, could
+    /// be quoted to convert it into a typing-only annotation.
+    ///
+    /// For example:
+    /// ```python
+    /// from pandas import DataFrame
+    ///
+    /// def foo() -> DataFrame:
+    ///    ...
+    /// ```
+    ///
+    /// Above, Python will evaluate `DataFrame` at runtime in order to add it to `__annotations__`.
+    RuntimeEvaluated,
+    /// Python will evaluate the annotation at runtime, and it's required to be available at
+    /// runtime, as a library (like Pydantic) needs access to it.
+    RuntimeRequired,
+    /// The annotation is only evaluated at type-checking time.
+    TypingOnly,
+}
+
+impl AnnotationContext {
+    pub(super) fn from_model(semantic: &SemanticModel, settings: &LinterSettings) -> Self {
+        // If the annotation is in a class scope (e.g., an annotated assignment for a
+        // class field), and that class is marked as annotation as runtime-required.
+        if semantic
+            .current_scope()
+            .kind
+            .as_class()
+            .is_some_and(|class_def| {
+                flake8_type_checking::helpers::runtime_required_class(
+                    class_def,
+                    &settings.flake8_type_checking.runtime_required_base_classes,
+                    &settings.flake8_type_checking.runtime_required_decorators,
+                    semantic,
+                )
+            })
+        {
+            return Self::RuntimeRequired;
+        }
+
+        // If `__future__` annotations are enabled, then annotations are never evaluated
+        // at runtime, so we can treat them as typing-only.
+        if semantic.future_annotations() {
+            return Self::TypingOnly;
+        }
+
+        // Otherwise, if we're in a class or module scope, then the annotation needs to
+        // be available at runtime.
+        // See: https://docs.python.org/3/reference/simple_stmts.html#annotated-assignment-statements
+        if matches!(
+            semantic.current_scope().kind,
+            ScopeKind::Class(_) | ScopeKind::Module
+        ) {
+            return Self::RuntimeEvaluated;
+        }
+
+        Self::TypingOnly
+    }
+}

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

@@ -58,6 +58,7 @@ use ruff_python_semantic::{
 use ruff_python_stdlib::builtins::{IPYTHON_BUILTINS, MAGIC_GLOBALS, PYTHON_BUILTINS};
 use ruff_source_file::Locator;
 
+use crate::checkers::ast::annotation::AnnotationContext;
 use crate::checkers::ast::deferred::Deferred;
 use crate::docstrings::extraction::ExtractionTarget;
 use crate::importer::Importer;
@@ -68,6 +69,7 @@ use crate::settings::{flags, LinterSettings};
 use crate::{docstrings, noqa};
 
 mod analyze;
+mod annotation;
 mod deferred;
 
 pub(crate) struct Checker<'a> {
@@ -679,62 +681,14 @@ where
                 value,
                 ..
             }) => {
-                enum AnnotationKind {
-                    RuntimeRequired,
-                    RuntimeEvaluated,
-                    TypingOnly,
-                }
-
-                fn annotation_kind(
-                    semantic: &SemanticModel,
-                    settings: &LinterSettings,
-                ) -> AnnotationKind {
-                    // If the annotation is in a class, and that class is marked as
-                    // runtime-evaluated, treat the annotation as runtime-required.
-                    // TODO(charlie): We could also include function calls here.
-                    if semantic
-                        .current_scope()
-                        .kind
-                        .as_class()
-                        .is_some_and(|class_def| {
-                            flake8_type_checking::helpers::runtime_required_class(
-                                class_def,
-                                &settings.flake8_type_checking.runtime_evaluated_base_classes,
-                                &settings.flake8_type_checking.runtime_evaluated_decorators,
-                                semantic,
-                            )
-                        })
-                    {
-                        return AnnotationKind::RuntimeRequired;
-                    }
-
-                    // If `__future__` annotations are enabled, then annotations are never evaluated
-                    // at runtime, so we can treat them as typing-only.
-                    if semantic.future_annotations() {
-                        return AnnotationKind::TypingOnly;
-                    }
-
-                    // Otherwise, if we're in a class or module scope, then the annotation needs to
-                    // be available at runtime.
-                    // See: https://docs.python.org/3/reference/simple_stmts.html#annotated-assignment-statements
-                    if matches!(
-                        semantic.current_scope().kind,
-                        ScopeKind::Class(_) | ScopeKind::Module
-                    ) {
-                        return AnnotationKind::RuntimeEvaluated;
-                    }
-
-                    AnnotationKind::TypingOnly
-                }
-
-                match annotation_kind(&self.semantic, self.settings) {
-                    AnnotationKind::RuntimeRequired => {
+                match AnnotationContext::from_model(&self.semantic, self.settings) {
+                    AnnotationContext::RuntimeRequired => {
                         self.visit_runtime_required_annotation(annotation);
                     }
-                    AnnotationKind::RuntimeEvaluated => {
+                    AnnotationContext::RuntimeEvaluated => {
                         self.visit_runtime_evaluated_annotation(annotation);
                     }
-                    AnnotationKind::TypingOnly => self.visit_annotation(annotation),
+                    AnnotationContext::TypingOnly => self.visit_annotation(annotation),
                 }
 
                 if let Some(expr) = value {

crates/ruff_linter/src/rules/flake8_type_checking/helpers.rs

@@ -6,12 +6,25 @@ use ruff_python_ast::call_path::from_qualified_name;
 use ruff_python_ast::helpers::{map_callable, map_subscript};
 use ruff_python_ast::{self as ast, Expr};
 use ruff_python_codegen::Stylist;
-use ruff_python_semantic::{Binding, BindingId, BindingKind, NodeId, SemanticModel};
+use ruff_python_semantic::{
+    Binding, BindingId, BindingKind, NodeId, ResolvedReference, SemanticModel,
+};
 use ruff_source_file::Locator;
 use ruff_text_size::Ranged;
 
 use crate::rules::flake8_type_checking::settings::Settings;
 
+/// Returns `true` if the [`ResolvedReference`] is in a typing-only context _or_ a runtime-evaluated
+/// context (with quoting enabled).
+pub(crate) fn is_typing_reference(reference: &ResolvedReference, settings: &Settings) -> bool {
+    reference.in_type_checking_block()
+        || reference.in_typing_only_annotation()
+        || reference.in_complex_string_type_definition()
+        || reference.in_simple_string_type_definition()
+        || (settings.quote_annotations && reference.in_runtime_evaluated_annotation())
+}
+
+/// Returns `true` if the [`Binding`] represents a runtime-required import.
 pub(crate) fn is_valid_runtime_import(
     binding: &Binding,
     semantic: &SemanticModel,
@@ -25,16 +38,7 @@ pub(crate) fn is_valid_runtime_import(
             && binding
                 .references()
                 .map(|reference_id| semantic.reference(reference_id))
-                .any(|reference| {
-                    // Are we in a typing-only context _or_ a runtime-evaluated context? We're
-                    // willing to quote runtime-evaluated, but not runtime-required annotations.
-                    !(reference.in_type_checking_block()
-                        || reference.in_typing_only_annotation()
-                        || reference.in_complex_string_type_definition()
-                        || reference.in_simple_string_type_definition()
-                        || (settings.quote_annotations
-                            && reference.in_runtime_evaluated_annotation()))
-                })
+                .any(|reference| !is_typing_reference(reference, settings))
     } else {
         false
     }
@@ -212,34 +216,34 @@ pub(crate) fn quote_annotation(
     locator: &Locator,
     stylist: &Stylist,
 ) -> Result<Edit> {
-    let expr = semantic.expression(node_id);
+    let expr = semantic.expression(node_id).expect("Expression not found");
     if let Some(parent_id) = semantic.parent_expression_id(node_id) {
         match semantic.expression(parent_id) {
-            Expr::Subscript(parent) => {
+            Some(Expr::Subscript(parent)) => {
                 if expr == parent.value.as_ref() {
                     // If we're quoting the value of a subscript, we need to quote the entire
                     // expression. For example, when quoting `DataFrame` in `DataFrame[int]`, we
                     // should generate `"DataFrame[int]"`.
                     return quote_annotation(parent_id, semantic, locator, stylist);
                 }
             }
-            Expr::Attribute(parent) => {
+            Some(Expr::Attribute(parent)) => {
                 if expr == parent.value.as_ref() {
                     // If we're quoting the value of an attribute, we need to quote the entire
                     // expression. For example, when quoting `DataFrame` in `pd.DataFrame`, we
                     // should generate `"pd.DataFrame"`.
                     return quote_annotation(parent_id, semantic, locator, stylist);
                 }
             }
-            Expr::Call(parent) => {
+            Some(Expr::Call(parent)) => {
                 if expr == parent.func.as_ref() {
                     // If we're quoting the function of a call, we need to quote the entire
                     // expression. For example, when quoting `DataFrame` in `DataFrame()`, we
                     // should generate `"DataFrame()"`.
                     return quote_annotation(parent_id, semantic, locator, stylist);
                 }
             }
-            Expr::BinOp(parent) => {
+            Some(Expr::BinOp(parent)) => {
                 if parent.op.is_bit_or() {
                     // If we're quoting the left or right side of a binary operation, we need to
                     // quote the entire expression. For example, when quoting `DataFrame` in
@@ -253,7 +257,12 @@ pub(crate) fn quote_annotation(
 
     let annotation = locator.slice(expr);
 
-    // If the annotation already contains a quote, avoid attempting to re-quote it.
+    // If the annotation already contains a quote, avoid attempting to re-quote it. For example:
+    // ```python
+    // from typing import Literal
+    //
+    // Set[Literal["Foo"]]
+    // ```
     if annotation.contains('\'') || annotation.contains('"') {
         return Err(anyhow::anyhow!("Annotation already contains a quote"));
     }

crates/ruff_linter/src/rules/flake8_type_checking/mod.rs

@@ -130,7 +130,7 @@ mod tests {
             Path::new("flake8_type_checking").join(path).as_path(),
             &settings::LinterSettings {
                 flake8_type_checking: super::settings::Settings {
-                    runtime_evaluated_base_classes: vec![
+                    runtime_required_base_classes: vec![
                         "pydantic.BaseModel".to_string(),
                         "sqlalchemy.orm.DeclarativeBase".to_string(),
                     ],
@@ -161,7 +161,7 @@ mod tests {
             Path::new("flake8_type_checking").join(path).as_path(),
             &settings::LinterSettings {
                 flake8_type_checking: super::settings::Settings {
-                    runtime_evaluated_decorators: vec![
+                    runtime_required_decorators: vec![
                         "attrs.define".to_string(),
                         "attrs.frozen".to_string(),
                     ],
@@ -186,7 +186,7 @@ mod tests {
             Path::new("flake8_type_checking").join(path).as_path(),
             &settings::LinterSettings {
                 flake8_type_checking: super::settings::Settings {
-                    runtime_evaluated_base_classes: vec!["module.direct.MyBaseClass".to_string()],
+                    runtime_required_base_classes: vec!["module.direct.MyBaseClass".to_string()],
                     ..Default::default()
                 },
                 ..settings::LinterSettings::for_rule(rule_code)

crates/ruff_linter/src/rules/flake8_type_checking/rules/runtime_import_in_type_checking_block.rs

@@ -72,7 +72,7 @@ impl Violation for RuntimeImportInTypeCheckingBlock {
                 "Move import `{qualified_name}` out of type-checking block. Import is used for more than type hinting."
             ),
             Strategy::QuoteUsages => format!(
-                "Quote references to `{qualified_name}`. Import is not available at runtime."
+                "Quote references to `{qualified_name}`. Import is in a type-checking block."
             ),
         }
     }
@@ -249,7 +249,7 @@ fn quote_imports(checker: &Checker, node_id: NodeId, imports: &[ImportBinding])
                 let reference = checker.semantic().reference(*reference_id);
                 if reference.context().is_runtime() {
                     Some(quote_annotation(
-                        reference.node_id()?,
+                        reference.expression_id()?,
                         checker.semantic(),
                         checker.locator(),
                         checker.stylist(),

crates/ruff_linter/src/rules/flake8_type_checking/rules/typing_only_runtime_import.rs

@@ -12,7 +12,7 @@ use crate::checkers::ast::Checker;
 use crate::codes::Rule;
 use crate::fix;
 use crate::importer::ImportedMembers;
-use crate::rules::flake8_type_checking::helpers::quote_annotation;
+use crate::rules::flake8_type_checking::helpers::{is_typing_reference, quote_annotation};
 use crate::rules::flake8_type_checking::imports::ImportBinding;
 use crate::rules::isort::{categorize, ImportSection, ImportType};
 
@@ -274,15 +274,7 @@ pub(crate) fn typing_only_runtime_import(
                 .references()
                 .map(|reference_id| checker.semantic().reference(reference_id))
                 .all(|reference| {
-                    // All references should be in a typing context _or_ a runtime-evaluated
-                    // annotation (as opposed to a runtime-required annotation), which we can
-                    // quote.
-                    reference.in_type_checking_block()
-                        || reference.in_typing_only_annotation()
-                        || reference.in_complex_string_type_definition()
-                        || reference.in_simple_string_type_definition()
-                        || (checker.settings.flake8_type_checking.quote_annotations
-                            && reference.in_runtime_evaluated_annotation())
+                    is_typing_reference(reference, &checker.settings.flake8_type_checking)
                 })
         {
             let qualified_name = import.qualified_name();
@@ -497,7 +489,7 @@ fn fix_imports(checker: &Checker, node_id: NodeId, imports: &[ImportBinding]) ->
                 let reference = checker.semantic().reference(*reference_id);
                 if reference.context().is_runtime() {
                     Some(quote_annotation(
-                        reference.node_id()?,
+                        reference.expression_id()?,
                         checker.semantic(),
                         checker.locator(),
                         checker.stylist(),

crates/ruff_linter/src/rules/flake8_type_checking/settings.rs

@@ -6,8 +6,8 @@ use ruff_macros::CacheKey;
 pub struct Settings {
     pub strict: bool,
     pub exempt_modules: Vec<String>,
-    pub runtime_evaluated_base_classes: Vec<String>,
-    pub runtime_evaluated_decorators: Vec<String>,
+    pub runtime_required_base_classes: Vec<String>,
+    pub runtime_required_decorators: Vec<String>,
     pub quote_annotations: bool,
 }
 
@@ -16,8 +16,8 @@ impl Default for Settings {
         Self {
             strict: false,
             exempt_modules: vec!["typing".to_string(), "typing_extensions".to_string()],
-            runtime_evaluated_base_classes: vec![],
-            runtime_evaluated_decorators: vec![],
+            runtime_required_base_classes: vec![],
+            runtime_required_decorators: vec![],
             quote_annotations: false,
         }
     }

crates/ruff_linter/src/rules/flake8_type_checking/snapshots/ruff_linter__rules__flake8_type_checking__tests__quote_runtime-import-in-type-checking-block_quote.py.snap

@@ -1,7 +1,7 @@
 ---
 source: crates/ruff_linter/src/rules/flake8_type_checking/mod.rs
 ---
-quote.py:64:28: TCH004 [*] Quote references to `pandas.DataFrame`. Import is not available at runtime.
+quote.py:64:28: TCH004 [*] Quote references to `pandas.DataFrame`. Import is in a type-checking block.
    |
 63 |     if TYPE_CHECKING:
 64 |         from pandas import DataFrame

crates/ruff_python_semantic/src/model.rs

@@ -995,11 +995,10 @@ impl<'a> SemanticModel<'a> {
 
     /// Return the [`Expr`] corresponding to the given [`NodeId`].
     #[inline]
-    pub fn expression(&self, node_id: NodeId) -> &'a Expr {
+    pub fn expression(&self, node_id: NodeId) -> Option<&'a Expr> {
         self.nodes
             .ancestor_ids(node_id)
             .find_map(|id| self.nodes[id].as_expression())
-            .expect("No expression found")
     }
 
     /// Given a [`Expr`], return its parent, if any.

crates/ruff_python_semantic/src/reference.rs

@@ -26,7 +26,7 @@ pub struct ResolvedReference {
 
 impl ResolvedReference {
     /// The expression that the reference occurs in.
-    pub const fn node_id(&self) -> Option<NodeId> {
+    pub const fn expression_id(&self) -> Option<NodeId> {
         self.node_id
     }