Only flip quotes if there is no Literal or Annotation

Signed-off-by: Shaygan <[email protected]>
astral-sh · Aug 23, 2024 · e8c0f90 · e8c0f90
1 parent 14f51e3
commit e8c0f90
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 5 deletions.
diff --git a/crates/ruff_linter/src/rules/flake8_type_checking/helpers.rs b/crates/ruff_linter/src/rules/flake8_type_checking/helpers.rs
@@ -1,4 +1,7 @@
 use anyhow::Result;
+use ast::str::Quote;
+use ast::visitor::{self, Visitor};
+use similar::DiffableStr;
 use std::cmp::Reverse;
 
 use ruff_diagnostics::Edit;
@@ -266,7 +269,18 @@ pub(crate) fn quote_annotation(
     let annotation = generator.expr(expr);
 
     let annotation_new = if annotation.contains(quote.as_char()) {
-        annotation.replace(quote.as_char(), &quote.opposite().as_char().to_string())
+        let mut quote_annotation = QuoteAnnotation {
+            can_remove: true,
+            annotation,
+        };
+        quote_annotation.visit_expr(expr);
+        if quote_annotation.can_remove {
+            quote_annotation.annotation.replace(quote.as_char(), "")
+        } else {
+            quote_annotation
+                .annotation
+                .replace(quote.as_char(), &quote.opposite().as_char().to_string())
+        }
     } else {
         annotation
     };
@@ -296,3 +310,31 @@ pub(crate) fn filter_contained(edits: Vec<Edit>) -> Vec<Edit> {
     }
     filtered
 }
+
+pub(crate) struct QuoteAnnotation {
+    can_remove: bool,
+    annotation: String,
+}
+
+impl<'a> visitor::Visitor<'a> for QuoteAnnotation {
+    fn visit_expr(&mut self, expr: &'a Expr) {
+        match expr {
+            Expr::Subscript(ast::ExprSubscript { value, slice, .. }) => {
+                if let Some(name) = value.as_name_expr() {
+                    if name.id.as_str() == "Literal" {
+                        self.can_remove = false;
+                    }
+                    if name.id.as_str() == "Annotation" {
+                        self.can_remove = false;
+                    }
+                }
+                visitor::walk_expr(self, expr);
+            }
+            // NOTE: check if string is inside literal or first parameter of annotation
+            Expr::StringLiteral(ast::ExprStringLiteral { value, .. }) => {
+                visitor::walk_expr(self, expr);
+            }
+            _ => visitor::walk_expr(self, expr),
+        }
+    }
+}
diff --git a/...r__rules__flake8_type_checking__tests__quote_typing-only-third-party-import_quote.py.snap b/...r__rules__flake8_type_checking__tests__quote_typing-only-third-party-import_quote.py.snap
@@ -80,7 +80,7 @@ quote.py:16:24: TCH002 [*] Move third-party import `pandas.DataFrame` into a typ
 16    |-    from pandas import DataFrame
 17 20 | 
 18    |-    def baz() -> DataFrame["int"]:
-   21 |+    def baz() -> "DataFrame['int']":
+   21 |+    def baz() -> "DataFrame[int]":
 19 22 |         ...
 20 23 | 
 21 24 | 
@@ -467,14 +467,14 @@ quote.py:102:44: TCH002 [*] Move third-party import `django.contrib.auth.models.
 102     |-    from django.contrib.auth.models import AbstractBaseUser
 103 106 | 
 104     |-    def foo(self, user: AbstractBaseUser["int"], view: "type[CondorBaseViewSet]"):
-    107 |+    def foo(self, user: "AbstractBaseUser['int']", view: "type[CondorBaseViewSet]"):
+    107 |+    def foo(self, user: "AbstractBaseUser[int]", view: "type[CondorBaseViewSet]"):
 105 108 |         pass
 106 109 | 
 107     |-    def foo(self, user: AbstractBaseUser['int']):
-    110 |+    def foo(self, user: "AbstractBaseUser['int']"):
+    110 |+    def foo(self, user: "AbstractBaseUser[int]"):
 108 111 |         pass
 109 112 | 
 110     |-    def foo(self, user: AbstractBaseUser['int', "str"]):
-    113 |+    def foo(self, user: "AbstractBaseUser['int', 'str']"):
+    113 |+    def foo(self, user: "AbstractBaseUser[int, str]"):
 111 114 |         pass
 112 115 |