Add autofix for magic methods (ANN204) (#3633)

LICENSE

@@ -393,7 +393,6 @@ are:
     THE SOFTWARE.
   """
 
-
 - autoflake, licensed as follows:
   """
     Copyright (C) 2012-2018 Steven Myint
@@ -417,6 +416,31 @@ are:
     SOFTWARE.
   """
 
+- autotyping, licensed as follows:
+  """
+    MIT License
+
+    Copyright (c) 2023 Jelle Zijlstra
+
+    Permission is hereby granted, free of charge, to any person obtaining a copy
+    of this software and associated documentation files (the "Software"), to deal
+    in the Software without restriction, including without limitation the rights
+    to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+    copies of the Software, and to permit persons to whom the Software is
+    furnished to do so, subject to the following conditions:
+
+    The above copyright notice and this permission notice shall be included in all
+    copies or substantial portions of the Software.
+
+    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+    IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+    FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+    AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+    LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+    OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+    SOFTWARE.
+  """
+
 - Flake8, licensed as follows:
   """
     == Flake8 License (MIT) ==

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

@@ -0,0 +1,42 @@
+class Foo:
+    def __str__(self):
+        ...
+
+    def __repr__(self):
+        ...
+
+    def __len__(self):
+        ...
+
+    def __length_hint__(self):
+        ...
+
+    def __init__(self):
+        ...
+
+    def __del__(self):
+        ...
+
+    def __bool__(self):
+        ...
+
+    def __bytes__(self):
+        ...
+
+    def __format__(self, format_spec):
+        ...
+
+    def __contains__(self, item):
+        ...
+
+    def __complex__(self):
+        ...
+
+    def __int__(self):
+        ...
+
+    def __float__(self):
+        ...
+
+    def __index__(self):
+        ...

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

@@ -7,7 +7,7 @@ use ruff_python_ast::source_code::Locator;
 use ruff_python_ast::types::Range;
 
 /// ANN204
-pub fn add_return_none_annotation(locator: &Locator, stmt: &Stmt) -> Result<Fix> {
+pub fn add_return_annotation(locator: &Locator, stmt: &Stmt, annotation: &str) -> Result<Fix> {
     let range = Range::from(stmt);
     let contents = locator.slice(range);
 
@@ -18,7 +18,7 @@ pub fn add_return_none_annotation(locator: &Locator, stmt: &Stmt) -> Result<Fix>
     for (start, tok, ..) in lexer::lex_located(contents, Mode::Module, range.location).flatten() {
         if seen_lpar && seen_rpar {
             if matches!(tok, Tok::Colon) {
-                return Ok(Fix::insertion(" -> None".to_string(), start));
+                return Ok(Fix::insertion(format!(" -> {annotation}"), start));
             }
         }
 

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

@@ -190,4 +190,14 @@ mod tests {
         assert_yaml_snapshot!(diagnostics);
         Ok(())
     }
+
+    #[test]
+    fn simple_magic_methods() -> Result<()> {
+        let diagnostics = test_path(
+            Path::new("flake8_annotations/simple_magic_methods.py"),
+            &Settings::for_rule(Rule::MissingReturnTypeSpecialMethod),
+        )?;
+        assert_yaml_snapshot!(diagnostics);
+        Ok(())
+    }
 }

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

@@ -5,10 +5,11 @@ use ruff_diagnostics::{AlwaysAutofixableViolation, Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, violation};
 use ruff_python_ast::helpers::ReturnStatementVisitor;
 use ruff_python_ast::types::Range;
-use ruff_python_ast::visibility;
 use ruff_python_ast::visibility::Visibility;
+use ruff_python_ast::visibility::{self};
 use ruff_python_ast::visitor::Visitor;
 use ruff_python_ast::{cast, helpers};
+use ruff_python_stdlib::typing::SIMPLE_MAGIC_RETURN_TYPES;
 
 use crate::checkers::ast::Checker;
 use crate::docstrings::definition::{Definition, DefinitionKind};
@@ -650,7 +651,7 @@ pub fn definition(
                         helpers::identifier_range(stmt, checker.locator),
                     ));
                 }
-            } else if is_method && visibility::is_init(cast::name(stmt)) {
+            } else if is_method && visibility::is_init(name) {
                 // Allow omission of return annotation in `__init__` functions, as long as at
                 // least one argument is typed.
                 if checker
@@ -667,7 +668,7 @@ pub fn definition(
                             helpers::identifier_range(stmt, checker.locator),
                         );
                         if checker.patch(diagnostic.kind.rule()) {
-                            match fixes::add_return_none_annotation(checker.locator, stmt) {
+                            match fixes::add_return_annotation(checker.locator, stmt, "None") {
                                 Ok(fix) => {
                                     diagnostic.amend(fix);
                                 }
@@ -677,18 +678,30 @@ pub fn definition(
                         diagnostics.push(diagnostic);
                     }
                 }
-            } else if is_method && visibility::is_magic(cast::name(stmt)) {
+            } else if is_method && visibility::is_magic(name) {
                 if checker
                     .settings
                     .rules
                     .enabled(Rule::MissingReturnTypeSpecialMethod)
                 {
-                    diagnostics.push(Diagnostic::new(
+                    let mut diagnostic = Diagnostic::new(
                         MissingReturnTypeSpecialMethod {
                             name: name.to_string(),
                         },
                         helpers::identifier_range(stmt, checker.locator),
-                    ));
+                    );
+                    let return_type = SIMPLE_MAGIC_RETURN_TYPES.get(name);
+                    if let Some(return_type) = return_type {
+                        if checker.patch(diagnostic.kind.rule()) {
+                            match fixes::add_return_annotation(checker.locator, stmt, return_type) {
+                                Ok(fix) => {
+                                    diagnostic.amend(fix);
+                                }
+                                Err(e) => error!("Failed to generate fix: {e}"),
+                            }
+                        }
+                    }
+                    diagnostics.push(diagnostic);
                 }
             } else {
                 match visibility {