Simplify ANN204 autofix to use Parameters range (#6793)

## Summary

This PR fixes the bug where the decorator parentheses weren't being considered
when computing the autofix for `ANN204`. The existing logic would only look
for balanced parentheses and not multiple pairs of parentheses.

The solution is to remove the logic to generate the autofix and use the
`Parameters` end range directly which includes the parentheses as well.

## Test Plan

Add test case for `ANN204` with decorator being called

fixes: #6790

crates/ruff/resources/test/fixtures/flake8_annotations/annotation_presence.py

@@ -152,3 +152,9 @@ def f(a: Union[str, bytes, Any]) -> None: ...
 def f(a: Optional[Any]) -> None: ...
 def f(a: Annotated[Any, ...]) -> None: ...
 def f(a: "Union[str, bytes, Any]") -> None: ...
+
+
+class Foo:
+    @decorator()
+    def __init__(self: "Foo", foo: int):
+       ...

crates/ruff/src/rules/flake8_annotations/fixes.rs

@@ -1,45 +0,0 @@
-use anyhow::{bail, Result};
-use ruff_python_ast::{PySourceType, Ranged};
-use ruff_python_parser::{lexer, AsMode, Tok};
-
-use ruff_diagnostics::Edit;
-use ruff_source_file::Locator;
-
-/// ANN204
-pub(crate) fn add_return_annotation<T: Ranged>(
-    statement: &T,
-    annotation: &str,
-    source_type: PySourceType,
-    locator: &Locator,
-) -> Result<Edit> {
-    let contents = &locator.contents()[statement.range()];
-
-    // Find the colon (following the `def` keyword).
-    let mut seen_lpar = false;
-    let mut seen_rpar = false;
-    let mut count = 0u32;
-    for (tok, range) in
-        lexer::lex_starts_at(contents, source_type.as_mode(), statement.start()).flatten()
-    {
-        if seen_lpar && seen_rpar {
-            if matches!(tok, Tok::Colon) {
-                return Ok(Edit::insertion(format!(" -> {annotation}"), range.start()));
-            }
-        }
-
-        if matches!(tok, Tok::Lpar) {
-            if count == 0 {
-                seen_lpar = true;
-            }
-            count = count.saturating_add(1);
-        }
-        if matches!(tok, Tok::Rpar) {
-            count = count.saturating_sub(1);
-            if count == 0 {
-                seen_rpar = true;
-            }
-        }
-    }
-
-    bail!("Unable to locate colon in function definition");
-}

crates/ruff/src/rules/flake8_annotations/mod.rs

@@ -1,5 +1,4 @@
 //! Rules from [flake8-annotations](https://pypi.org/project/flake8-annotations/).
-mod fixes;
 pub(crate) mod helpers;
 pub(crate) mod rules;
 pub mod settings;

crates/ruff/src/rules/flake8_annotations/rules/definition.rs

@@ -1,4 +1,4 @@
-use ruff_diagnostics::{AlwaysAutofixableViolation, Diagnostic, Fix, Violation};
+use ruff_diagnostics::{AlwaysAutofixableViolation, Diagnostic, Edit, Fix, Violation};
 use ruff_macros::{derive_message_formats, violation};
 use ruff_python_ast::helpers::ReturnStatementVisitor;
 use ruff_python_ast::identifier::Identifier;
@@ -11,7 +11,6 @@ use ruff_python_stdlib::typing::simple_magic_return_type;
 
 use crate::checkers::ast::Checker;
 use crate::registry::{AsRule, Rule};
-use crate::rules::flake8_annotations::fixes;
 use crate::rules::ruff::typing::type_hint_resolves_to_any;
 
 /// ## What it does
@@ -704,15 +703,10 @@ pub(crate) fn definition(
                         function.identifier(),
                     );
                     if checker.patch(diagnostic.kind.rule()) {
-                        diagnostic.try_set_fix(|| {
-                            fixes::add_return_annotation(
-                                function,
-                                "None",
-                                checker.source_type,
-                                checker.locator(),
-                            )
-                            .map(Fix::suggested)
-                        });
+                        diagnostic.set_fix(Fix::suggested(Edit::insertion(
+                            " -> None".to_string(),
+                            function.parameters.range().end(),
+                        )));
                     }
                     diagnostics.push(diagnostic);
                 }
@@ -727,15 +721,10 @@ pub(crate) fn definition(
                 );
                 if checker.patch(diagnostic.kind.rule()) {
                     if let Some(return_type) = simple_magic_return_type(name) {
-                        diagnostic.try_set_fix(|| {
-                            fixes::add_return_annotation(
-                                function,
-                                return_type,
-                                checker.source_type,
-                                checker.locator(),
-                            )
-                            .map(Fix::suggested)
-                        });
+                        diagnostic.set_fix(Fix::suggested(Edit::insertion(
+                            format!(" -> {return_type}"),
+                            function.parameters.range().end(),
+                        )));
                     }
                 }
                 diagnostics.push(diagnostic);

crates/ruff/src/rules/flake8_annotations/snapshots/ruff__rules__flake8_annotations__tests__defaults.snap

@@ -242,4 +242,22 @@ annotation_presence.py:154:10: ANN401 Dynamically typed expressions (typing.Any)
     |          ^^^^^^^^^^^^^^^^^^^^^^^^ ANN401
     |
 
+annotation_presence.py:159:9: ANN204 [*] Missing return type annotation for special method `__init__`
+    |
+157 | class Foo:
+158 |     @decorator()
+159 |     def __init__(self: "Foo", foo: int):
+    |         ^^^^^^^^ ANN204
+160 |        ...
+    |
+    = help: Add `None` return type
+
+ℹ Suggested fix
+156 156 | 
+157 157 | class Foo:
+158 158 |     @decorator()
+159     |-    def __init__(self: "Foo", foo: int):
+    159 |+    def __init__(self: "Foo", foo: int) -> None:
+160 160 |        ...
+