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

This commit is contained in:
Jonathan Plasse 2023-03-21 00:19:20 +01:00 committed by GitHub
parent f039bf36a2
commit f70a49ed8b
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
7 changed files with 405 additions and 10 deletions

View file

@ -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):
...

View file

@ -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));
}
}

View file

@ -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(())
}
}

View file

@ -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 {

View file

@ -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: ~

View file

@ -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"),
])
});