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

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

@@ -1,3 +1,6 @@
+from typing import Literal
+
+
 def f():
     from pandas import DataFrame
 
@@ -97,3 +100,83 @@ def f():
 
     def func(self) -> DataFrame | list[Series]:
         pass
+
+def f():
+    from django.contrib.auth.models import AbstractBaseUser
+    from typing import Annotated
+
+    def test_remove_inner_quotes_double(self, user: AbstractBaseUser["int"], view: "type[CondorBaseViewSet]"):
+        pass
+
+    def test_remove_inner_quotes_single(self, user: AbstractBaseUser['int']):
+        pass
+
+    def test_remove_inner_quotes_mixed(self, user: AbstractBaseUser['int', "str"]):
+        pass
+
+    def test_literal_annotation_args_contain_quotes(self, type1: AbstractBaseUser[Literal["user", "admin"], Annotated["int", "1", 2]]):
+        pass
+
+    def test_union_contain_inner_quotes(self, type1: AbstractBaseUser["int" | Literal["int"]]):
+        pass
+
+    def test_literal_cannot_remove(user: AbstractBaseUser[Literal['user', "admin"]]):
+        pass
+
+    def test_literal_cannot_remove2(user: AbstractBaseUser[Literal['int'], str]):
+        pass
+
+    def test_type_literal_mixed_quotes(view: type["CondorBaseViewSet"], role: Literal["admin", 'user']):
+        pass
+
+    def test_mixed_quotes_literal(user: AbstractBaseUser[Literal['user'], "int"]):
+        pass
+
+    def test_annotated_literal_mixed_quotes(user: Annotated[str, "max_length=20", Literal['user', "admin"]]):
+        pass
+
+
+    def test_type_annotated_mixed(view: type["CondorBaseViewSet"], data: Annotated[int, "min_value=1", Literal["admin", 'superuser']]):
+        pass
+
+    def test_union_literal_mixed_quotes(user: Union[Literal['active', "inactive"], str]):
+        pass
+
+    def test_callable_literal_mixed_quotes(callable_fn: Callable[["int", Literal['admin', "user"]], 'bool']):
+        pass
+
+    def test_callable_annotated_literal(callable_fn: Callable[[int, Annotated[str, Literal['active', "inactive"]]], bool]):
+        pass
+
+    def test_generic_alias_literal_mixed_quotes(user: Union["List[int]", Literal['admin', "user"], 'Dict[str, Any]']):
+        pass
+
+    def test_literal_no_inner_removal(arg: Literal["user", "admin"]):
+        pass
+
+    def test_quoted_annotation_literal(arg: "Literal['user', 'admin']"):
+        pass
+
+    def test_annotated_with_expression(arg: Annotated[str, "min_length=3"]):
+        pass
+
+    def test_union_literal_no_removal(arg: Union[Literal["true", "false"], None]):
+        pass
+
+    def test_optional_literal_no_removal(arg: Optional[Literal["red", "blue"]]):
+        pass
+
+    def test_type_alias_with_removal(arg: "List[Dict[str, int]]"):
+        pass
+
+    def test_final_literal_no_removal(arg: Final[Literal["open", "close"]]):
+        pass
+
+    def test_callable_literal_no_removal(arg: Callable[[], Literal["success", "failure"]]):
+        pass
+
+    def test_classvar_union_literal_no_removal(arg: ClassVar[Union[Literal["yes", "no"], None]]):
+        pass
+
+    def test_tuple_literal_no_removal(arg: tuple[Literal["apple", "orange"], int]):
+        pass

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

@@ -1,4 +1,7 @@
 use anyhow::Result;
+use ast::str::Quote;
+use ast::visitor::source_order;
+use ruff_python_ast::visitor::source_order::SourceOrderVisitor;
 use std::cmp::Reverse;
 
 use ruff_diagnostics::Edit;
@@ -9,7 +12,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;
@@ -221,9 +223,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) {
@@ -233,51 +233,42 @@ 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_annotation = QuoteAnnotation::new(stylist);
+    quote_annotation.visit_expr(expr);
+
+    let annotation = quote_annotation.annotation;
 
     Ok(Edit::range_replacement(
         format!("{quote}{annotation}{quote}"),
@@ -304,3 +295,105 @@ pub(crate) fn filter_contained(edits: Vec<Edit>) -> Vec<Edit> {
     }
     filtered
 }
+
+#[derive(Copy, PartialEq, Clone)]
+enum State {
+    Literal,
+    AnnotatedFirst,
+    AnnotatedRest,
+    Other,
+}
+
+pub(crate) struct QuoteAnnotation<'a> {
+    state: Vec<State>,
+    stylist: &'a Stylist<'a>,
+    annotation: String,
+    final_quote_type: Quote,
+}
+
+impl<'a> QuoteAnnotation<'a> {
+    pub(crate) fn new(stylist: &'a Stylist<'a>) -> Self {
+        let final_quote_type = stylist.quote();
+        Self {
+            state: vec![],
+            stylist,
+            annotation: String::new(),
+            final_quote_type,
+        }
+    }
+}
+
+impl<'a> source_order::SourceOrderVisitor<'a> for QuoteAnnotation<'a> {
+    fn visit_expr(&mut self, expr: &'a Expr) {
+        let generator = Generator::from(self.stylist);
+        match expr {
+            Expr::Subscript(ast::ExprSubscript { value, slice, .. }) => {
+                if let Some(name) = value.as_name_expr() {
+                    let value = generator.expr(value);
+                    self.annotation.push_str(&value);
+                    self.annotation.push('[');
+                    match name.id.as_str() {
+                        "Literal" => self.state.push(State::Literal),
+                        "Annotated" => self.state.push(State::AnnotatedFirst),
+                        _ => self.state.push(State::Other),
+                    }
+
+                    self.visit_expr(slice);
+                    self.state.pop();
+                    self.annotation.push(']');
+                }
+            }
+            Expr::Tuple(ast::ExprTuple { elts, .. }) => {
+                let Some(first_elm) = elts.first() else {
+                    return;
+                };
+                self.visit_expr(first_elm);
+                if self.state.last().copied() == Some(State::AnnotatedFirst) {
+                    self.state.push(State::AnnotatedRest);
+                }
+                for elm in elts.iter().skip(1) {
+                    self.annotation.push_str(", ");
+                    self.visit_expr(elm);
+                }
+                self.state.pop();
+            }
+            Expr::BinOp(ast::ExprBinOp {
+                left, op, right, ..
+            }) => {
+                self.visit_expr(left);
+                self.annotation.push(' ');
+                self.annotation.push_str(op.as_str());
+                self.annotation.push(' ');
+                self.visit_expr(right);
+            }
+            Expr::BoolOp(ast::ExprBoolOp { op, values, .. }) => {
+                for (i, value) in values.iter().enumerate() {
+                    if i > 0 {
+                        self.annotation.push(' ');
+                        self.annotation.push_str(op.as_str());
+                    }
+                    self.visit_expr(value);
+                }
+            }
+            _ => {
+                let source = match self.state.last().copied() {
+                    Some(State::Literal | State::AnnotatedRest) => {
+                        let mut source = generator.expr(expr);
+                        source = source.replace(
+                            self.final_quote_type.as_char(),
+                            &self.final_quote_type.opposite().as_char().to_string(),
+                        );
+                        source
+                    }
+                    None | Some(State::AnnotatedFirst | State::Other) => {
+                        let mut source = generator.expr(expr);
+                        source = source.replace(self.final_quote_type.as_char(), "");
+                        source = source.replace(self.final_quote_type.opposite().as_char(), "");
+                        source
+                    }
+                };
+                self.annotation.push_str(&source);
+            }
+        }
+    }
+}

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

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,24 +1,22 @@
 ---
 source: crates/ruff_linter/src/rules/flake8_type_checking/mod.rs
 ---
-quote.py:64:28: TCH004 [*] Quote references to `pandas.DataFrame`. Import is in a type-checking block.
+quote.py:67:28: TCH004 [*] Quote references to `pandas.DataFrame`. Import is in a type-checking block.
    |
-63 |     if TYPE_CHECKING:
-64 |         from pandas import DataFrame
+66 |     if TYPE_CHECKING:
+67 |         from pandas import DataFrame
    |                            ^^^^^^^^^ TCH004
-65 | 
-66 |     def func(value: DataFrame):
+68 | 
+69 |     def func(value: DataFrame):
    |
    = help: Quote references
 
 ℹ Unsafe fix
-63 63 |     if TYPE_CHECKING:
-64 64 |         from pandas import DataFrame
-65 65 | 
-66    |-    def func(value: DataFrame):
-   66 |+    def func(value: "DataFrame"):
-67 67 |         ...
+66 66 |     if TYPE_CHECKING:
+67 67 |         from pandas import DataFrame
 68 68 | 
-69 69 | 
-
-
+69    |-    def func(value: DataFrame):
+   69 |+    def func(value: "DataFrame"):
+70 70 |         ...
+71 71 | 
+72 72 |