[pyupgrade] Handle comments and multiline expressions correctly (`U…

…P037`) (#15337)

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

@@ -0,0 +1,53 @@
+# 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''']): ...
+
+
+def f(a: Foo['''
+Bar[
+    Multi |
+    Line]  # Comment''']): ...
+
+
+def f(a: '''list[int]
+	''' = []): ...
+
+
+a: '''\\
+list[int]''' = [42]
+
+
+# TODO: These are valid too. String annotations are assumed to be enclosed in parentheses.
+# https://github.com/astral-sh/ruff/issues/9467
+
+def f(a: '''
+	list[int]
+	''' = []): ...
+
+
+def f(a: Foo['''
+    Bar
+    [
+    Multi |
+    Line
+    ]  # Comment''']): ...
+
+
+a: '''list
+[int]''' = [42]

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,12 @@
-use ruff_text_size::TextRange;
+use ruff_text_size::{TextLen, 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_ast::Stmt;
+use ruff_python_semantic::SemanticModel;
+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.
@@ -79,10 +82,38 @@ 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 placeholder_range = TextRange::up_to(annotation.text_len());
+    let spans_multiple_lines = annotation.contains_line_break(placeholder_range);
+
+    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) if in_parameter_annotation(range.start(), checker.semantic()) => {
+            annotation.to_string()
+        }
+        (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));
+}
+
+fn in_parameter_annotation(offset: TextSize, semantic: &SemanticModel) -> bool {
+    let Stmt::FunctionDef(stmt) = semantic.current_statement() else {
+        return false;
+    };
+
+    stmt.parameters.range.contains(offset)
 }

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

@@ -0,0 +1,162 @@
+---
+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 |+)]): ...
+20 21 | 
+21 22 | 
+22 23 | def f(a: Foo['''
+
+UP037_2.pyi:22:14: UP037 [*] Remove quotes from type annotation
+   |
+22 |   def f(a: Foo['''
+   |  ______________^
+23 | | Bar[
+24 | |     Multi |
+25 | |     Line]  # Comment''']): ...
+   | |_______________________^ UP037
+   |
+   = help: Remove quotes
+
+ℹ Safe fix
+19 19 |     Line]  # Comment''']): ...
+20 20 | 
+21 21 | 
+22    |-def f(a: Foo['''
+   22 |+def f(a: Foo[(
+23 23 | Bar[
+24 24 |     Multi |
+25    |-    Line]  # Comment''']): ...
+   25 |+    Line]  # Comment
+   26 |+)]): ...
+26 27 | 
+27 28 | 
+28 29 | def f(a: '''list[int]
+
+UP037_2.pyi:28:10: UP037 [*] Remove quotes from type annotation
+   |
+28 |   def f(a: '''list[int]
+   |  __________^
+29 | |     ''' = []): ...
+   | |_______^ UP037
+   |
+   = help: Remove quotes
+
+ℹ Safe fix
+25 25 |     Line]  # Comment''']): ...
+26 26 | 
+27 27 | 
+28    |-def f(a: '''list[int]
+29    |-	''' = []): ...
+   28 |+def f(a: list[int]
+   29 |+	 = []): ...
+30 30 | 
+31 31 | 
+32 32 | a: '''\\
+
+UP037_2.pyi:32:4: UP037 [*] Remove quotes from type annotation
+   |
+32 |   a: '''\\
+   |  ____^
+33 | | list[int]''' = [42]
+   | |____________^ UP037
+   |
+   = help: Remove quotes
+
+ℹ Safe fix
+29 29 | 	''' = []): ...
+30 30 | 
+31 31 | 
+32    |-a: '''\\
+33    |-list[int]''' = [42]
+   32 |+a: (\
+   33 |+list[int]) = [42]
+34 34 | 
+35 35 | 
+36 36 | # TODO: These are valid too. String annotations are assumed to be enclosed in parentheses.