Remove unnecessary quotes in annotation expressions

Signed-off-by: Shaygan <hey@glyphack.com>

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

@@ -1,8 +1,121 @@
+from typing import Literal
+
+
+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:
+        ...
+
+
+def f():
+    import pandas as pd
+
+    def baz() -> pd.DataFrame.Extra:
+        ...
+
+
+def f():
+    import pandas as pd
+
+    def baz() -> pd.DataFrame | int:
+        ...
+
+
+
+def f():
+    from pandas import DataFrame
+
+    def baz() -> DataFrame():
+        ...
+
+
+def f():
+    from typing import Literal
+
+    from pandas import DataFrame
+
+    def baz() -> DataFrame[Literal["int"]]:
+        ...
+
+
+def f():
+    from typing import TYPE_CHECKING
+
+    if TYPE_CHECKING:
+        from pandas import DataFrame
+
+    def func(value: DataFrame):
+        ...
+
+
+def f():
+    from pandas import DataFrame, Series
+
+    def baz() -> DataFrame | Series:
+        ...
+
+
+def f():
+    from pandas import DataFrame, Series
+
+    def baz() -> (
+        DataFrame |
+        Series
+    ):
+        ...
+
+    class C:
+        x: DataFrame[
+            int
+        ] = 1
+
+    def func() -> DataFrame[[DataFrame[_P, _R]], DataFrame[_P, _R]]:
+        ...
+
+
+def f():
+    from pandas import DataFrame, Series
+
+    def func(self) -> DataFrame | list[Series]:
+        pass
+
 def f():
     from django.contrib.auth.models import AbstractBaseUser
+    from typing import Annotated
 
-    def x (
-    user: AbstractBaseUser["int"],
-   ):
+    def foo(self, user: AbstractBaseUser["int"], view: "type[CondorBaseViewSet]"):
         pass
 
+    def foo(self, user: AbstractBaseUser['int']):
+        pass
+
+    def foo(self, user: AbstractBaseUser['int', "str"]):
+        pass
+
+    def set_role(self, type1: AbstractBaseUser[Literal["user", "admin"], Annotated["int", "1", 2]]):
+        pass
+
+    def test_role(self, type1: AbstractBaseUser["int" | Literal["int"]]):
+        pass

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

@@ -1,8 +1,7 @@
 use anyhow::Result;
 use ast::str::Quote;
-use ast::visitor::transformer::Transformer;
-use ast::visitor::{self, transformer, Visitor};
-use ast::{ExprStringLiteral, StringLiteralFlags, StringLiteralValue};
+use ast::visitor::source_order;
+use ruff_python_ast::visitor::source_order::SourceOrderVisitor;
 use std::cmp::Reverse;
 
 use ruff_diagnostics::Edit;
@@ -267,20 +266,12 @@ pub(crate) fn quote_annotation(
     }
 
     let quote = stylist.quote();
-    let mut quote_annotation = QuoteAnnotation {
-        can_remove: vec![],
-        generator: &generator,
-        stylist: &stylist,
-        contain_single_quote: false,
-        contain_double_quote: false,
-        annotation: String::new(),
-    };
+    let mut quote_annotation = QuoteAnnotation::new(stylist);
+    quote_annotation.visit_expr(&expr);
 
-    let mut new_expr = expr.clone();
-    quote_annotation.visit_expr(&mut new_expr);
-
-    let annotation = generator.expr(&new_expr);
+    let annotation = quote_annotation.annotation;
 
+    dbg!(&annotation);
     Ok(Edit::range_replacement(
         format!("{quote}{annotation}{quote}"),
         expr.range(),
@@ -307,40 +298,90 @@ pub(crate) fn filter_contained(edits: Vec<Edit>) -> Vec<Edit> {
     filtered
 }
 
+#[derive(Copy, PartialEq, Clone)]
+enum State {
+    Literal,
+    Annotated,
+    AnnotatedNonFirstElm,
+    Other,
+}
+
 pub(crate) struct QuoteAnnotation<'a> {
-    can_remove: Vec<bool>,
-    generator: &'a Generator<'a>,
+    state: Vec<State>,
     stylist: &'a Stylist<'a>,
-    contain_single_quote: bool,
-    contain_double_quote: bool,
     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> transformer::Transformer for QuoteAnnotation<'a> {
-    fn visit_expr(&self, expr: &mut Expr) {
+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() {
-                    if name.id.as_str() == "Literal" {
-                        todo!()
-                    }
-                    if name.id.as_str() == "Annotation" {
-                        todo!()
+                    let value = generator.expr(value);
+                    self.annotation.push_str(&format!("{value}["));
+                    match name.id.as_str() {
+                        "Literal" => self.state.push(State::Literal),
+                        "Annotated" => self.state.push(State::Annotated),
+                        _ => self.state.push(State::Other),
                     }
-                    transformer::walk_expr(self, expr);
+
+                    self.visit_expr(slice);
+                    self.state.pop();
+                    self.annotation.push_str(&format!("]"));
                 }
             }
-            Expr::StringLiteral(ast::ExprStringLiteral { value, range }) => {
-                let new_str = value.to_string();
-                let node = Expr::from(ast::ExprName {
-                    id: new_str,
-                    range: *range,
-                    ctx: ast::ExprContext::Load,
-                });
-                *expr = node;
-                transformer::walk_expr(self, expr);
+            Expr::Tuple(ast::ExprTuple { elts, .. }) => {
+                let first_elm = elts.first().unwrap();
+                self.visit_expr(first_elm);
+                if self.state.last().copied() == Some(State::Annotated) {
+                    self.state.push(State::AnnotatedNonFirstElm);
+                }
+                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_str(&format!(" {op} "));
+                self.visit_expr(right);
+            }
+            _ => {
+                let source = match self.state.last().copied() {
+                    Some(State::Literal | State::Annotated) => {
+                        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
+                    }
+                    _ => {
+                        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);
             }
-            _ => transformer::walk_expr(self, expr),
         }
     }
 }

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 |