[pyupgrade] Handle comments and multiline expressions correctly (UP037)

crates/ruff_linter/resources/test/fixtures/pyupgrade/UP037_2.pyi

@@ -0,0 +1,19 @@
+# https://github.com/astral-sh/ruff/issues/7102
+
+def f(a: Foo['SingleLine  # Comment']): ...
+
+
+def f(a: Foo['''Bar[
+    Multi |
+    Line]''']): ...
+
+
+def f(a: Foo['''Bar[
+    Multi |
+    Line  # Comment
+]''']): ...
+
+
+def f(a: Foo['''Bar[
+    Multi |
+    Line]  # Comment''']): ...

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

@@ -58,6 +58,7 @@ mod tests {
     #[test_case(Rule::PrintfStringFormatting, Path::new("UP031_1.py"))]
     #[test_case(Rule::QuotedAnnotation, Path::new("UP037_0.py"))]
     #[test_case(Rule::QuotedAnnotation, Path::new("UP037_1.py"))]
+    #[test_case(Rule::QuotedAnnotation, Path::new("UP037_2.pyi"))]
     #[test_case(Rule::RedundantOpenModes, Path::new("UP015.py"))]
     #[test_case(Rule::RedundantOpenModes, Path::new("UP015_1.py"))]
     #[test_case(Rule::ReplaceStdoutStderr, Path::new("UP022.py"))]

crates/ruff_linter/src/rules/pyupgrade/rules/quoted_annotation.rs

@@ -1,9 +1,10 @@
-use ruff_text_size::TextRange;
+use ruff_text_size::{TextRange, TextSize};
 
+use crate::checkers::ast::Checker;
 use ruff_diagnostics::{AlwaysFixableViolation, Diagnostic, Edit, Fix};
 use ruff_macros::{derive_message_formats, ViolationMetadata};
-
-use crate::checkers::ast::Checker;
+use ruff_python_trivia::{SimpleToken, SimpleTokenKind, SimpleTokenizer};
+use ruff_source_file::LineRanges;
 
 /// ## What it does
 /// Checks for the presence of unnecessary quotes in type annotations.
@@ -73,10 +74,28 @@ impl AlwaysFixableViolation for QuotedAnnotation {
 
 /// UP037
 pub(crate) fn quoted_annotation(checker: &mut Checker, annotation: &str, range: TextRange) {
-    let mut diagnostic = Diagnostic::new(QuotedAnnotation, range);
-    diagnostic.set_fix(Fix::safe_edit(Edit::range_replacement(
-        annotation.to_string(),
-        range,
-    )));
-    checker.diagnostics.push(diagnostic);
+    let diagnostic = Diagnostic::new(QuotedAnnotation, range);
+
+    let len = TextSize::try_from(annotation.len()).unwrap();
+    let placeholder_range = TextRange::up_to(len);
+    let spans_multiple_lines = annotation.count_lines(placeholder_range) > 1;
+
+    let tokenizer = SimpleTokenizer::new(annotation, placeholder_range);
+    let last_token_is_comment = matches!(
+        tokenizer.last(),
+        Some(SimpleToken {
+            kind: SimpleTokenKind::Comment,
+            ..
+        })
+    );
+
+    let new_content = match (spans_multiple_lines, last_token_is_comment) {
+        (false, false) => annotation.to_string(),
+        (true, false) => format!("({annotation})"),
+        (_, true) => format!("({annotation}\n)"),
+    };
+    let edit = Edit::range_replacement(new_content, range);
+    let fix = Fix::safe_edit(edit);
+
+    checker.diagnostics.push(diagnostic.with_fix(fix));
 }

crates/ruff_linter/src/rules/pyupgrade/snapshots/ruff_linter__rules__pyupgrade__tests__UP037_2.pyi.snap

@@ -0,0 +1,91 @@
+---
+source: crates/ruff_linter/src/rules/pyupgrade/mod.rs
+snapshot_kind: text
+---
+UP037_2.pyi:3:14: UP037 [*] Remove quotes from type annotation
+  |
+1 | # https://github.com/astral-sh/ruff/issues/7102
+2 | 
+3 | def f(a: Foo['SingleLine  # Comment']): ...
+  |              ^^^^^^^^^^^^^^^^^^^^^^^ UP037
+  |
+  = help: Remove quotes
+
+ℹ Safe fix
+1 1 | # https://github.com/astral-sh/ruff/issues/7102
+2 2 | 
+3   |-def f(a: Foo['SingleLine  # Comment']): ...
+  3 |+def f(a: Foo[(SingleLine  # Comment
+  4 |+)]): ...
+4 5 | 
+5 6 | 
+6 7 | def f(a: Foo['''Bar[
+
+UP037_2.pyi:6:14: UP037 [*] Remove quotes from type annotation
+  |
+6 |   def f(a: Foo['''Bar[
+  |  ______________^
+7 | |     Multi |
+8 | |     Line]''']): ...
+  | |____________^ UP037
+  |
+  = help: Remove quotes
+
+ℹ Safe fix
+3 3 | def f(a: Foo['SingleLine  # Comment']): ...
+4 4 | 
+5 5 | 
+6   |-def f(a: Foo['''Bar[
+  6 |+def f(a: Foo[(Bar[
+7 7 |     Multi |
+8   |-    Line]''']): ...
+  8 |+    Line])]): ...
+9 9 | 
+10 10 | 
+11 11 | def f(a: Foo['''Bar[
+
+UP037_2.pyi:11:14: UP037 [*] Remove quotes from type annotation
+   |
+11 |   def f(a: Foo['''Bar[
+   |  ______________^
+12 | |     Multi |
+13 | |     Line  # Comment
+14 | | ]''']): ...
+   | |____^ UP037
+   |
+   = help: Remove quotes
+
+ℹ Safe fix
+8  8  |     Line]''']): ...
+9  9  | 
+10 10 | 
+11    |-def f(a: Foo['''Bar[
+   11 |+def f(a: Foo[(Bar[
+12 12 |     Multi |
+13 13 |     Line  # Comment
+14    |-]''']): ...
+   14 |+])]): ...
+15 15 | 
+16 16 | 
+17 17 | def f(a: Foo['''Bar[
+
+UP037_2.pyi:17:14: UP037 [*] Remove quotes from type annotation
+   |
+17 |   def f(a: Foo['''Bar[
+   |  ______________^
+18 | |     Multi |
+19 | |     Line]  # Comment''']): ...
+   | |_______________________^ UP037
+   |
+   = help: Remove quotes
+
+ℹ Safe fix
+14 14 | ]''']): ...
+15 15 | 
+16 16 | 
+17    |-def f(a: Foo['''Bar[
+   17 |+def f(a: Foo[(Bar[
+18 18 |     Multi |
+19    |-    Line]  # Comment''']): ...
+   19 |+    Line]  # Comment
+   20 |+)]): ...