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 {

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

@@ -0,0 +1,285 @@
+---
+source: crates/ruff/src/rules/flake8_annotations/mod.rs
+expression: diagnostics
+---
+- kind:
+    name: MissingReturnTypeSpecialMethod
+    body: "Missing return type annotation for special method `__str__`"
+    suggestion: "Add `None` return type"
+    fixable: true
+  location:
+    row: 2
+    column: 8
+  end_location:
+    row: 2
+    column: 15
+  fix:
+    content: " -> str"
+    location:
+      row: 2
+      column: 21
+    end_location:
+      row: 2
+      column: 21
+  parent: ~
+- kind:
+    name: MissingReturnTypeSpecialMethod
+    body: "Missing return type annotation for special method `__repr__`"
+    suggestion: "Add `None` return type"
+    fixable: true
+  location:
+    row: 5
+    column: 8
+  end_location:
+    row: 5
+    column: 16
+  fix:
+    content: " -> str"
+    location:
+      row: 5
+      column: 22
+    end_location:
+      row: 5
+      column: 22
+  parent: ~
+- kind:
+    name: MissingReturnTypeSpecialMethod
+    body: "Missing return type annotation for special method `__len__`"
+    suggestion: "Add `None` return type"
+    fixable: true
+  location:
+    row: 8
+    column: 8
+  end_location:
+    row: 8
+    column: 15
+  fix:
+    content: " -> int"
+    location:
+      row: 8
+      column: 21
+    end_location:
+      row: 8
+      column: 21
+  parent: ~
+- kind:
+    name: MissingReturnTypeSpecialMethod
+    body: "Missing return type annotation for special method `__length_hint__`"
+    suggestion: "Add `None` return type"
+    fixable: true
+  location:
+    row: 11
+    column: 8
+  end_location:
+    row: 11
+    column: 23
+  fix:
+    content: " -> int"
+    location:
+      row: 11
+      column: 29
+    end_location:
+      row: 11
+      column: 29
+  parent: ~
+- kind:
+    name: MissingReturnTypeSpecialMethod
+    body: "Missing return type annotation for special method `__init__`"
+    suggestion: "Add `None` return type"
+    fixable: true
+  location:
+    row: 14
+    column: 8
+  end_location:
+    row: 14
+    column: 16
+  fix:
+    content: " -> None"
+    location:
+      row: 14
+      column: 22
+    end_location:
+      row: 14
+      column: 22
+  parent: ~
+- kind:
+    name: MissingReturnTypeSpecialMethod
+    body: "Missing return type annotation for special method `__del__`"
+    suggestion: "Add `None` return type"
+    fixable: true
+  location:
+    row: 17
+    column: 8
+  end_location:
+    row: 17
+    column: 15
+  fix:
+    content: " -> None"
+    location:
+      row: 17
+      column: 21
+    end_location:
+      row: 17
+      column: 21
+  parent: ~
+- kind:
+    name: MissingReturnTypeSpecialMethod
+    body: "Missing return type annotation for special method `__bool__`"
+    suggestion: "Add `None` return type"
+    fixable: true
+  location:
+    row: 20
+    column: 8
+  end_location:
+    row: 20
+    column: 16
+  fix:
+    content: " -> bool"
+    location:
+      row: 20
+      column: 22
+    end_location:
+      row: 20
+      column: 22
+  parent: ~
+- kind:
+    name: MissingReturnTypeSpecialMethod
+    body: "Missing return type annotation for special method `__bytes__`"
+    suggestion: "Add `None` return type"
+    fixable: true
+  location:
+    row: 23
+    column: 8
+  end_location:
+    row: 23
+    column: 17
+  fix:
+    content: " -> bytes"
+    location:
+      row: 23
+      column: 23
+    end_location:
+      row: 23
+      column: 23
+  parent: ~
+- kind:
+    name: MissingReturnTypeSpecialMethod
+    body: "Missing return type annotation for special method `__format__`"
+    suggestion: "Add `None` return type"
+    fixable: true
+  location:
+    row: 26
+    column: 8
+  end_location:
+    row: 26
+    column: 18
+  fix:
+    content: " -> str"
+    location:
+      row: 26
+      column: 37
+    end_location:
+      row: 26
+      column: 37
+  parent: ~
+- kind:
+    name: MissingReturnTypeSpecialMethod
+    body: "Missing return type annotation for special method `__contains__`"
+    suggestion: "Add `None` return type"
+    fixable: true
+  location:
+    row: 29
+    column: 8
+  end_location:
+    row: 29
+    column: 20
+  fix:
+    content: " -> bool"
+    location:
+      row: 29
+      column: 32
+    end_location:
+      row: 29
+      column: 32
+  parent: ~
+- kind:
+    name: MissingReturnTypeSpecialMethod
+    body: "Missing return type annotation for special method `__complex__`"
+    suggestion: "Add `None` return type"
+    fixable: true
+  location:
+    row: 32
+    column: 8
+  end_location:
+    row: 32
+    column: 19
+  fix:
+    content: " -> complex"
+    location:
+      row: 32
+      column: 25
+    end_location:
+      row: 32
+      column: 25
+  parent: ~
+- kind:
+    name: MissingReturnTypeSpecialMethod
+    body: "Missing return type annotation for special method `__int__`"
+    suggestion: "Add `None` return type"
+    fixable: true
+  location:
+    row: 35
+    column: 8
+  end_location:
+    row: 35
+    column: 15
+  fix:
+    content: " -> int"
+    location:
+      row: 35
+      column: 21
+    end_location:
+      row: 35
+      column: 21
+  parent: ~
+- kind:
+    name: MissingReturnTypeSpecialMethod
+    body: "Missing return type annotation for special method `__float__`"
+    suggestion: "Add `None` return type"
+    fixable: true
+  location:
+    row: 38
+    column: 8
+  end_location:
+    row: 38
+    column: 17
+  fix:
+    content: " -> float"
+    location:
+      row: 38
+      column: 23
+    end_location:
+      row: 38
+      column: 23
+  parent: ~
+- kind:
+    name: MissingReturnTypeSpecialMethod
+    body: "Missing return type annotation for special method `__index__`"
+    suggestion: "Add `None` return type"
+    fixable: true
+  location:
+    row: 41
+    column: 8
+  end_location:
+    row: 41
+    column: 17
+  fix:
+    content: " -> int"
+    location:
+      row: 41
+      column: 23
+    end_location:
+      row: 41
+      column: 23
+  parent: ~
+

crates/ruff_python_stdlib/src/typing.rs

@@ -1,5 +1,5 @@
 use once_cell::sync::Lazy;
-use rustc_hash::FxHashSet;
+use rustc_hash::{FxHashMap, FxHashSet};
 
 // See: https://pypi.org/project/typing-extensions/
 pub static TYPING_EXTENSIONS: Lazy<FxHashSet<&'static str>> = Lazy::new(|| {
@@ -197,3 +197,24 @@ pub const PEP_585_BUILTINS_ELIGIBLE: &[&[&str]] = &[
     &["typing", "Type"],
     &["typing_extensions", "Type"],
 ];
+
+// See: https://github.com/JelleZijlstra/autotyping/blob/f65b5ee3a8fdb77999f84b4c87edb996e25269a5/autotyping/autotyping.py#L69-L84
+pub static SIMPLE_MAGIC_RETURN_TYPES: Lazy<FxHashMap<&'static str, &'static str>> =
+    Lazy::new(|| {
+        FxHashMap::from_iter([
+            ("__str__", "str"),
+            ("__repr__", "str"),
+            ("__len__", "int"),
+            ("__length_hint__", "int"),
+            ("__init__", "None"),
+            ("__del__", "None"),
+            ("__bool__", "bool"),
+            ("__bytes__", "bytes"),
+            ("__format__", "str"),
+            ("__contains__", "bool"),
+            ("__complex__", "complex"),
+            ("__int__", "int"),
+            ("__float__", "float"),
+            ("__index__", "int"),
+        ])
+    });