Review feedback

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

@@ -689,9 +689,8 @@ where
                     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 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

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/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/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
     }