[flake8-type-checking] Support auto-quoting when annotations contain quotes

crates/ruff_linter/resources/test/fixtures/flake8_type_checking/quote.py

@@ -12,13 +12,6 @@ def baz() -> DataFrame[int]:
         ...
 
 
-def f():
-    from pandas import DataFrame
-
-    def baz() -> DataFrame["int"]:
-        ...
-
-
 def f():
     import pandas as pd
 

crates/ruff_linter/resources/test/fixtures/flake8_type_checking/quote2.py

@@ -0,0 +1,87 @@
+def f():
+    from django.contrib.auth.models import AbstractBaseUser
+
+    def test_remove_inner_quotes_double(self, user: AbstractBaseUser["int"]):
+        pass
+
+
+def f():
+    from django.contrib.auth.models import AbstractBaseUser
+
+    def test_remove_inner_quotes_single(self, user: AbstractBaseUser['int']):
+        pass
+
+
+def f():
+    from django.contrib.auth.models import AbstractBaseUser
+
+    def test_remove_inner_quotes_mixed(self, user: AbstractBaseUser['int', "str"]):
+        pass
+
+
+def f():
+    from typing import Annotated, Literal
+
+    from django.contrib.auth.models import AbstractBaseUser
+
+    def test_literal_annotation_args_contain_quotes(self, type1: AbstractBaseUser[Literal["user", "admin"], Annotated["int", "1", 2]]):
+        pass
+
+
+def f():
+    from typing import Literal
+
+    from django.contrib.auth.models import AbstractBaseUser
+
+    def test_union_contain_inner_quotes(self, type1: AbstractBaseUser["int" | Literal["int"]]):
+        pass
+
+
+def f():
+    from typing import Literal
+
+    from django.contrib.auth.models import AbstractBaseUser
+
+    def test_inner_literal_mixed_quotes(user: AbstractBaseUser[Literal['user', "admin"]]):
+        pass
+
+
+def f():
+    from typing import Literal
+
+    from django.contrib.auth.models import AbstractBaseUser
+
+    def test_inner_literal_single_quote(user: AbstractBaseUser[Literal['int'], str]):
+        pass
+
+
+def f():
+    from typing import Literal
+
+    from django.contrib.auth.models import AbstractBaseUser
+
+    def test_mixed_quotes_literal(user: AbstractBaseUser[Literal['user'], "int"]):
+        pass
+
+
+def f():
+    from typing import Annotated, Literal
+
+    from django.contrib.auth.models import AbstractBaseUser
+
+    def test_annotated_literal_mixed_quotes(user: AbstractBaseUser[Annotated[str, "max_length=20", Literal['user', "admin"]]]):
+        pass
+
+
+def f():
+    from typing import Literal, Union
+
+    def test_union_literal_no_removal(arg: Union[Literal["true", "false"], None]):
+        pass
+
+
+def f():
+    from typing import Literal, Optional
+
+    def test_optional_literal_no_removal(arg: Optional[Literal["red", "blue"]]):
+        pass

crates/ruff_linter/resources/test/fixtures/flake8_type_checking/quote3.py

@@ -0,0 +1,44 @@
+def f():
+    from typing import Literal, Union
+
+    from django.contrib.auth.models import AbstractBaseUser
+
+    def test_union_literal_mixed_quotes(user: AbstractBaseUser[Union[Literal['active', "inactive"], str]]):
+        pass
+
+
+def f():
+    from typing import Callable, Literal
+
+    from django.contrib.auth.models import AbstractBaseUser
+
+    def test_callable_literal_mixed_quotes(callable_fn: AbstractBaseUser[Callable[["int", Literal['admin', "user"]], 'bool']]):
+        pass
+
+
+def f():
+    from typing import Annotated, Callable, Literal
+
+    from django.contrib.auth.models import AbstractBaseUser
+
+    def test_callable_annotated_literal(callable_fn: AbstractBaseUser[Callable[[int, Annotated[str, Literal['active', "inactive"]]], bool]]):
+        pass
+
+
+def f():
+    from typing import literal
+
+    from django.contrib.auth import models
+
+    def test_attribute(arg: models.AbstractBaseUser["int"]):
+        pass
+
+
+def f():
+    from typing import Literal
+
+    from django.contrib.auth import models
+
+    def test_attribute_typing_literal(arg: models.AbstractBaseUser[Literal["admin"]]):
+        pass
+

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

@@ -1,4 +1,6 @@
 use anyhow::Result;
+use ast::visitor::source_order;
+use ruff_python_ast::visitor::source_order::SourceOrderVisitor;
 use std::cmp::Reverse;
 
 use ruff_diagnostics::Edit;
@@ -9,7 +11,6 @@ use ruff_python_codegen::{Generator, Stylist};
 use ruff_python_semantic::{
     analyze, Binding, BindingKind, Modules, NodeId, ResolvedReference, ScopeKind, SemanticModel,
 };
-use ruff_source_file::Locator;
 use ruff_text_size::Ranged;
 
 use crate::rules::flake8_type_checking::settings::Settings;
@@ -223,9 +224,7 @@ pub(crate) fn is_singledispatch_implementation(
 pub(crate) fn quote_annotation(
     node_id: NodeId,
     semantic: &SemanticModel,
-    locator: &Locator,
     stylist: &Stylist,
-    generator: Generator,
 ) -> Result<Edit> {
     let expr = semantic.expression(node_id).expect("Expression not found");
     if let Some(parent_id) = semantic.parent_expression_id(node_id) {
@@ -235,51 +234,41 @@ pub(crate) fn quote_annotation(
                     // 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, generator);
+                    return quote_annotation(parent_id, semantic, stylist);
                 }
             }
             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, generator);
+                    return quote_annotation(parent_id, semantic, stylist);
                 }
             }
             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, generator);
+                    return quote_annotation(parent_id, semantic, stylist);
                 }
             }
             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
                     // `DataFrame | Series`, we should generate `"DataFrame | Series"`.
-                    return quote_annotation(parent_id, semantic, locator, stylist, generator);
+                    return quote_annotation(parent_id, semantic, stylist);
                 }
             }
             _ => {}
         }
     }
 
-    // If the annotation already contains a quote, avoid attempting to re-quote it. For example:
-    // ```python
-    // from typing import Literal
-    //
-    // Set[Literal["Foo"]]
-    // ```
-    let text = locator.slice(expr);
-    if text.contains('\'') || text.contains('"') {
-        return Err(anyhow::anyhow!("Annotation already contains a quote"));
-    }
-
-    // Quote the entire expression.
     let quote = stylist.quote();
-    let annotation = generator.expr(expr);
+    let mut quote_annotator = QuoteAnnotator::new(semantic, stylist);
+    quote_annotator.visit_expr(expr);
+    let annotation = quote_annotator.into_annotation();
 
     Ok(Edit::range_replacement(
         format!("{quote}{annotation}{quote}"),
@@ -306,3 +295,111 @@ pub(crate) fn filter_contained(edits: Vec<Edit>) -> Vec<Edit> {
     }
     filtered
 }
+
+#[derive(Copy, PartialEq, Clone)]
+enum QuoteAnnotatorState {
+    Literal,
+    AnnotatedFirst,
+    AnnotatedRest,
+    Other,
+}
+
+pub(crate) struct QuoteAnnotator<'a> {
+    stylist: &'a Stylist<'a>,
+    semantic: &'a SemanticModel<'a>,
+    state: Vec<QuoteAnnotatorState>,
+    annotation: String,
+}
+
+impl<'a> QuoteAnnotator<'a> {
+    fn new(semantic: &'a SemanticModel<'a>, stylist: &'a Stylist<'a>) -> Self {
+        Self {
+            stylist,
+            semantic,
+            state: Vec::new(),
+            annotation: String::new(),
+        }
+    }
+
+    fn into_annotation(self) -> String {
+        self.annotation
+    }
+}
+
+impl<'a> source_order::SourceOrderVisitor<'a> for QuoteAnnotator<'a> {
+    fn visit_expr(&mut self, expr: &'a Expr) {
+        let generator = Generator::from(self.stylist);
+
+        match expr {
+            Expr::Subscript(ast::ExprSubscript { value, slice, .. }) => {
+                let value_expr = generator.expr(value);
+                self.annotation.push_str(&value_expr);
+                self.annotation.push('[');
+
+                if let Some(qualified_name) = self.semantic.resolve_qualified_name(value) {
+                    if self
+                        .semantic
+                        .match_typing_qualified_name(&qualified_name, "Literal")
+                    {
+                        self.state.push(QuoteAnnotatorState::Literal);
+                    } else if self
+                        .semantic
+                        .match_typing_qualified_name(&qualified_name, "Annotated")
+                    {
+                        self.state.push(QuoteAnnotatorState::AnnotatedFirst);
+                    } else {
+                        self.state.push(QuoteAnnotatorState::Other);
+                    }
+                }
+
+                self.visit_expr(slice);
+                self.state.pop();
+                self.annotation.push(']');
+            }
+            Expr::Tuple(ast::ExprTuple { elts, .. }) => {
+                let Some((first, remaining)) = elts.split_first() else {
+                    return;
+                };
+                self.visit_expr(first);
+                if let Some(last) = self.state.last_mut() {
+                    if *last == QuoteAnnotatorState::AnnotatedFirst {
+                        *last = QuoteAnnotatorState::AnnotatedRest;
+                    }
+                }
+                for expr in remaining {
+                    self.annotation.push_str(", ");
+                    self.visit_expr(expr);
+                }
+                self.state.pop();
+            }
+            Expr::BinOp(ast::ExprBinOp {
+                left, op, right, ..
+            }) => {
+                debug_assert_eq!(*op, ast::Operator::BitOr);
+                self.visit_expr(left);
+                self.annotation.push_str(" | ");
+                self.visit_expr(right);
+            }
+            _ => {
+                let source = match self.state.last().copied() {
+                    Some(QuoteAnnotatorState::Literal | QuoteAnnotatorState::AnnotatedRest) => {
+                        let mut source = generator.expr(expr);
+                        source = source.replace(
+                            self.stylist.quote().as_char(),
+                            &self.stylist.quote().opposite().as_char().to_string(),
+                        );
+                        source
+                    }
+                    None
+                    | Some(QuoteAnnotatorState::AnnotatedFirst | QuoteAnnotatorState::Other) => {
+                        let mut source = generator.expr(expr);
+                        source = source.replace(self.stylist.quote().as_char(), "");
+                        source = source.replace(self.stylist.quote().opposite().as_char(), "");
+                        source
+                    }
+                };
+                self.annotation.push_str(&source);
+            }
+        }
+    }
+}

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

@@ -60,6 +60,10 @@ mod tests {
 
     #[test_case(Rule::RuntimeImportInTypeCheckingBlock, Path::new("quote.py"))]
     #[test_case(Rule::TypingOnlyThirdPartyImport, Path::new("quote.py"))]
+    #[test_case(Rule::RuntimeImportInTypeCheckingBlock, Path::new("quote2.py"))]
+    #[test_case(Rule::TypingOnlyThirdPartyImport, Path::new("quote2.py"))]
+    #[test_case(Rule::RuntimeImportInTypeCheckingBlock, Path::new("quote3.py"))]
+    #[test_case(Rule::TypingOnlyThirdPartyImport, Path::new("quote3.py"))]
     fn quote(rule_code: Rule, path: &Path) -> Result<()> {
         let snapshot = format!("quote_{}_{}", rule_code.as_ref(), path.to_string_lossy());
         let diagnostics = test_path(

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

@@ -270,9 +270,7 @@ fn quote_imports(checker: &Checker, node_id: NodeId, imports: &[ImportBinding])
                         Some(quote_annotation(
                             reference.expression_id()?,
                             checker.semantic(),
-                            checker.locator(),
                             checker.stylist(),
-                            checker.generator(),
                         ))
                     } else {
                         None

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

@@ -505,9 +505,7 @@ fn fix_imports(checker: &Checker, node_id: NodeId, imports: &[ImportBinding]) ->
                         Some(quote_annotation(
                             reference.expression_id()?,
                             checker.semantic(),
-                            checker.locator(),
                             checker.stylist(),
-                            checker.generator(),
                         ))
                     } else {
                         None