Add quote support

crates/ruff/resources/test/fixtures/flake8_type_checking/TCH002.py

@@ -172,3 +172,31 @@ def f():
     from module import Member
 
     x: Member = 1
+
+
+def f():
+    from pandas import DataFrame
+
+    def baz() -> DataFrame:
+        ...
+
+
+def f():
+    from pandas import DataFrame
+
+    def baz() -> DataFrame[int]:
+        ...
+
+
+def f():
+    from pandas import DataFrame
+
+    def baz() -> DataFrame["int"]:
+        ...
+
+
+def f():
+    import pandas as pd
+
+    def baz() -> pd.DataFrame:
+        ...

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

@@ -1891,7 +1891,7 @@ where
                 {
                     if let Some(expr) = &arg_with_default.def.annotation {
                         if runtime_annotation {
-                            self.visit_runtime_annotation(expr);
+                            self.visit_runtime_evaluated_annotation(expr);
                         } else {
                             self.visit_annotation(expr);
                         };
@@ -1903,7 +1903,7 @@ where
                 if let Some(arg) = &args.vararg {
                     if let Some(expr) = &arg.annotation {
                         if runtime_annotation {
-                            self.visit_runtime_annotation(expr);
+                            self.visit_runtime_evaluated_annotation(expr);
                         } else {
                             self.visit_annotation(expr);
                         };
@@ -1912,15 +1912,15 @@ where
                 if let Some(arg) = &args.kwarg {
                     if let Some(expr) = &arg.annotation {
                         if runtime_annotation {
-                            self.visit_runtime_annotation(expr);
+                            self.visit_runtime_evaluated_annotation(expr);
                         } else {
                             self.visit_annotation(expr);
                         };
                     }
                 }
                 for expr in returns {
                     if runtime_annotation {
-                        self.visit_runtime_annotation(expr);
+                        self.visit_runtime_evaluated_annotation(expr);
                     } else {
                         self.visit_annotation(expr);
                     };
@@ -2032,38 +2032,54 @@ where
                 value,
                 ..
             }) => {
-                // 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
-                let runtime_annotation = if self.semantic.future_annotations() {
-                    if self.semantic.scope().kind.is_class() {
-                        let baseclasses = &self
-                            .settings
-                            .flake8_type_checking
-                            .runtime_evaluated_base_classes;
-                        let decorators = &self
-                            .settings
-                            .flake8_type_checking
-                            .runtime_evaluated_decorators;
-                        flake8_type_checking::helpers::runtime_evaluated(
+                enum AnnotationKind {
+                    RuntimeRequired,
+                    RuntimeEvaluated,
+                    TypingOnly,
+                }
+
+                fn annotation_kind(model: &SemanticModel, settings: &Settings) -> 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 model.scope().kind.is_class() {
+                        let baseclasses =
+                            &settings.flake8_type_checking.runtime_evaluated_base_classes;
+                        let decorators =
+                            &settings.flake8_type_checking.runtime_evaluated_decorators;
+                        if flake8_type_checking::helpers::runtime_required(
                             baseclasses,
                             decorators,
-                            &self.semantic,
-                        )
-                    } else {
-                        false
+                            model,
+                        ) {
+                            return AnnotationKind::RuntimeRequired;
+                        }
                     }
-                } else {
-                    matches!(
-                        self.semantic.scope().kind,
-                        ScopeKind::Class(_) | ScopeKind::Module
-                    )
-                };
 
-                if runtime_annotation {
-                    self.visit_runtime_annotation(annotation);
-                } else {
-                    self.visit_annotation(annotation);
+                    // If `__future__` annotations are enabled, then annotations are never evaluated
+                    // at runtime, so we can treat them as typing-only.
+                    if model.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!(model.scope().kind, ScopeKind::Class(_) | ScopeKind::Module) {
+                        return AnnotationKind::RuntimeEvaluated;
+                    }
+
+                    AnnotationKind::TypingOnly
+                }
+
+                match annotation_kind(&self.semantic, self.settings) {
+                    AnnotationKind::RuntimeRequired => {
+                        self.visit_runtime_required_annotation(annotation);
+                    }
+                    AnnotationKind::RuntimeEvaluated => {
+                        self.visit_runtime_evaluated_annotation(annotation);
+                    }
+                    AnnotationKind::TypingOnly => self.visit_annotation(annotation),
                 }
                 if let Some(expr) = value {
                     if self.semantic.match_typing_expr(annotation, "TypeAlias") {
@@ -4304,10 +4320,18 @@ impl<'a> Checker<'a> {
         self.semantic.flags = snapshot;
     }
 
+    /// Visit an [`Expr`], and treat it as a runtime-evaluated type annotation.
+    fn visit_runtime_evaluated_annotation(&mut self, expr: &'a Expr) {
+        let snapshot = self.semantic.flags;
+        self.semantic.flags |= SemanticModelFlags::RUNTIME_EVALUATED_ANNOTATION;
+        self.visit_type_definition(expr);
+        self.semantic.flags = snapshot;
+    }
+
     /// Visit an [`Expr`], and treat it as a runtime-required type annotation.
-    fn visit_runtime_annotation(&mut self, expr: &'a Expr) {
+    fn visit_runtime_required_annotation(&mut self, expr: &'a Expr) {
         let snapshot = self.semantic.flags;
-        self.semantic.flags |= SemanticModelFlags::RUNTIME_ANNOTATION;
+        self.semantic.flags |= SemanticModelFlags::RUNTIME_REQUIRED_ANNOTATION;
         self.visit_type_definition(expr);
         self.semantic.flags = snapshot;
     }

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

@@ -12,31 +12,40 @@ pub(crate) fn is_valid_runtime_import(binding: &Binding, semantic: &SemanticMode
         binding.context.is_runtime()
             && binding
                 .references()
-                .any(|reference_id| semantic.reference(reference_id).context().is_runtime())
+                .map(|reference_id| semantic.reference(reference_id))
+                .any(|reference| {
+                    // This is like: typing context _or_ a runtime-required type annotation (since
+                    // we're willing to quote it).
+                    !(reference.in_type_checking_block()
+                        || reference.in_typing_only_annotation()
+                        || reference.in_runtime_evaluated_annotation()
+                        || reference.in_complex_string_type_definition()
+                        || reference.in_simple_string_type_definition())
+                })
     } else {
         false
     }
 }
 
-pub(crate) fn runtime_evaluated(
+pub(crate) fn runtime_required(
     base_classes: &[String],
     decorators: &[String],
     semantic: &SemanticModel,
 ) -> bool {
     if !base_classes.is_empty() {
-        if runtime_evaluated_base_class(base_classes, semantic) {
+        if runtime_required_base_class(base_classes, semantic) {
             return true;
         }
     }
     if !decorators.is_empty() {
-        if runtime_evaluated_decorators(decorators, semantic) {
+        if runtime_required_decorators(decorators, semantic) {
             return true;
         }
     }
     false
 }
 
-fn runtime_evaluated_base_class(base_classes: &[String], semantic: &SemanticModel) -> bool {
+fn runtime_required_base_class(base_classes: &[String], semantic: &SemanticModel) -> bool {
     if let ScopeKind::Class(ast::StmtClassDef { bases, .. }) = &semantic.scope().kind {
         for base in bases {
             if let Some(call_path) = semantic.resolve_call_path(base) {
@@ -52,7 +61,7 @@ fn runtime_evaluated_base_class(base_classes: &[String], semantic: &SemanticMode
     false
 }
 
-fn runtime_evaluated_decorators(decorators: &[String], semantic: &SemanticModel) -> bool {
+fn runtime_required_decorators(decorators: &[String], semantic: &SemanticModel) -> bool {
     if let ScopeKind::Class(ast::StmtClassDef { decorator_list, .. }) = &semantic.scope().kind {
         for decorator in decorator_list {
             if let Some(call_path) = semantic.resolve_call_path(map_callable(&decorator.expression))

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

@@ -1,9 +1,10 @@
 use anyhow::Result;
-use ruff_text_size::TextRange;
+use ruff_text_size::{TextLen, TextRange, TextSize};
 use rustc_hash::FxHashMap;
 
-use ruff_diagnostics::{AutofixKind, Diagnostic, DiagnosticKind, Fix, Violation};
+use ruff_diagnostics::{AutofixKind, Diagnostic, DiagnosticKind, Edit, Fix, Violation};
 use ruff_macros::{derive_message_formats, violation};
+use ruff_python_ast::source_code::{Locator, Stylist};
 use ruff_python_semantic::{Binding, NodeId, ResolvedReferenceId, Scope};
 
 use crate::autofix;
@@ -228,13 +229,19 @@ pub(crate) fn typing_only_runtime_import(
         };
 
         if binding.context.is_runtime()
-            && binding.references().all(|reference_id| {
-                checker
-                    .semantic()
-                    .reference(reference_id)
-                    .context()
-                    .is_typing()
-            })
+            && binding
+                .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_runtime_evaluated_annotation()
+                        || reference.in_complex_string_type_definition()
+                        || reference.in_simple_string_type_definition()
+                })
         {
             // Extract the module base and level from the full name.
             // Ex) `foo.bar.baz` -> `foo`, `0`
@@ -278,6 +285,7 @@ pub(crate) fn typing_only_runtime_import(
             let import = Import {
                 qualified_name,
                 reference_id,
+                binding,
                 range: binding.range,
                 parent_range: binding.parent_range(checker.semantic()),
             };
@@ -356,6 +364,8 @@ pub(crate) fn typing_only_runtime_import(
 struct Import<'a> {
     /// The qualified name of the import (e.g., `typing.List` for `from typing import List`).
     qualified_name: &'a str,
+    /// The binding for the imported symbol.
+    binding: &'a Binding<'a>,
     /// The first reference to the imported symbol.
     reference_id: ResolvedReferenceId,
     /// The trimmed range of the import (e.g., `List` in `from typing import List`).
@@ -449,8 +459,98 @@ fn fix_imports(checker: &Checker, stmt_id: NodeId, imports: &[Import]) -> Result
         checker.semantic(),
     )?;
 
-    Ok(
-        Fix::suggested_edits(remove_import_edit, add_import_edit.into_edits())
-            .isolate(checker.isolation(parent)),
+    // Step 3) Quote any runtime usages of the referenced symbol.
+    let quote_reference_edits = imports.iter().flat_map(|Import { binding, .. }| {
+        binding.references.iter().filter_map(|reference_id| {
+            let reference = checker.semantic().reference(*reference_id);
+            if reference.context().is_runtime() {
+                Some(quote_annotation(
+                    reference.range(),
+                    checker.locator(),
+                    checker.stylist(),
+                ))
+            } else {
+                None
+            }
+        })
+    });
+
+    Ok(Fix::suggested_edits(
+        remove_import_edit,
+        add_import_edit
+            .into_edits()
+            .into_iter()
+            .chain(quote_reference_edits),
     )
+    .isolate(checker.isolation(parent)))
+}
+
+/// Quote a type annotation.
+///
+/// This requires more than wrapping the reference in quotes. For example:
+/// - When quoting `Series` in `Series[pd.Timestamp]`, we want `"Series[pd.Timestamp]"`.
+/// - When quoting `kubernetes` in `kubernetes.SecurityContext`, we want `"kubernetes.SecurityContext"`.
+/// - When quoting `Series` in `Series["pd.Timestamp"]`, we want `"Series[pd.Timestamp]"`.
+fn quote_annotation(range: TextRange, locator: &Locator, stylist: &Stylist) -> Edit {
+    // Expand the annotation to the end of the reference.
+    let mut depth = 0u32;
+    let mut len = TextSize::default();
+    let mut annotation = String::with_capacity(range.len().into());
+    for c in locator.after(range.start()).chars() {
+        match c {
+            '[' => depth += 1,
+            ']' => {
+                // Ex) Quoting `int` in `DataFrame[int]`, which should expand until the end of the
+                // `int` symbol`.
+                if depth == 0 {
+                    break;
+                }
+
+                depth -= 1;
+
+                // Ex) Quoting `DataFrame` in `DataFrame[int]`, which should expand until the end
+                // of the subscript.
+                if depth == 0 {
+                    annotation.push(c);
+                    len += c.text_len();
+                    break;
+                }
+            }
+            '.' => {
+                // Expand attributes.
+            }
+            'a'..='z' | 'A'..='Z' | '_' | '0'..='9' => {
+                // Expand identifiers.
+            }
+            '"' | '\'' => {
+                // Skip quotes.
+                // TODO(charlie): Retain escaped quotes, and quotes in literals.
+                len += c.text_len();
+                continue;
+            }
+            '\n' | '\r' if depth > 0 => {
+                // If we hit a newline, fallback to replacing the range. This can be ugly, but is
+                // better than not quoting at all.
+                let annotation = locator.slice(range);
+                let quote = stylist.quote();
+                let annotation = format!("{quote}{annotation}{quote}");
+                return Edit::range_replacement(annotation, range);
+            }
+            _ => {
+                // If we hit a space, or a parenthesis, or any other character (and we're not in
+                // a subscript), we're done.
+                if depth == 0 {
+                    break;
+                }
+            }
+        }
+        annotation.push(c);
+        len += c.text_len();
+    }
+
+    // Wrap in quotes.
+    let quote = stylist.quote();
+    let annotation = format!("{quote}{annotation}{quote}");
+
+    Edit::range_replacement(annotation, TextRange::at(range.start(), len))
 }