From f9c8f7ee61f3653122ad98e8202c334748dd3210 Mon Sep 17 00:00:00 2001 From: Dhruv Manilawala Date: Wed, 23 Aug 2023 01:28:58 +0530 Subject: [PATCH 1/2] Consider decorator parenthese for `ANN204` autofix --- .../flake8_annotations/annotation_presence.py | 6 ++++++ .../ruff/src/rules/flake8_annotations/fixes.rs | 1 + ...s__flake8_annotations__tests__defaults.snap | 18 ++++++++++++++++++ 3 files changed, 25 insertions(+) diff --git a/crates/ruff/resources/test/fixtures/flake8_annotations/annotation_presence.py b/crates/ruff/resources/test/fixtures/flake8_annotations/annotation_presence.py index 6c512262687ad..36c5155b83e79 100644 --- a/crates/ruff/resources/test/fixtures/flake8_annotations/annotation_presence.py +++ b/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): + ... diff --git a/crates/ruff/src/rules/flake8_annotations/fixes.rs b/crates/ruff/src/rules/flake8_annotations/fixes.rs index 1608732475ba1..4d1a880956cca 100644 --- a/crates/ruff/src/rules/flake8_annotations/fixes.rs +++ b/crates/ruff/src/rules/flake8_annotations/fixes.rs @@ -30,6 +30,7 @@ pub(crate) fn add_return_annotation( if matches!(tok, Tok::Lpar) { if count == 0 { seen_lpar = true; + seen_rpar = false; } count = count.saturating_add(1); } diff --git a/crates/ruff/src/rules/flake8_annotations/snapshots/ruff__rules__flake8_annotations__tests__defaults.snap b/crates/ruff/src/rules/flake8_annotations/snapshots/ruff__rules__flake8_annotations__tests__defaults.snap index 54e30dd686e64..03944c6d78df5 100644 --- a/crates/ruff/src/rules/flake8_annotations/snapshots/ruff__rules__flake8_annotations__tests__defaults.snap +++ b/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 | ... + From 53bc842da63853da2b53e1594329f47512a7ddbd Mon Sep 17 00:00:00 2001 From: Dhruv Manilawala Date: Wed, 23 Aug 2023 08:54:29 +0530 Subject: [PATCH 2/2] Simplify `ANN204` autofix to use `Parameter` range --- .../src/rules/flake8_annotations/fixes.rs | 46 ------------------- .../ruff/src/rules/flake8_annotations/mod.rs | 1 - .../flake8_annotations/rules/definition.rs | 29 ++++-------- 3 files changed, 9 insertions(+), 67 deletions(-) delete mode 100644 crates/ruff/src/rules/flake8_annotations/fixes.rs diff --git a/crates/ruff/src/rules/flake8_annotations/fixes.rs b/crates/ruff/src/rules/flake8_annotations/fixes.rs deleted file mode 100644 index 4d1a880956cca..0000000000000 --- a/crates/ruff/src/rules/flake8_annotations/fixes.rs +++ /dev/null @@ -1,46 +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( - statement: &T, - annotation: &str, - source_type: PySourceType, - locator: &Locator, -) -> Result { - 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; - seen_rpar = false; - } - 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"); -} diff --git a/crates/ruff/src/rules/flake8_annotations/mod.rs b/crates/ruff/src/rules/flake8_annotations/mod.rs index 9928a48fb1434..2322045c272b0 100644 --- a/crates/ruff/src/rules/flake8_annotations/mod.rs +++ b/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; diff --git a/crates/ruff/src/rules/flake8_annotations/rules/definition.rs b/crates/ruff/src/rules/flake8_annotations/rules/definition.rs index ae750a3f23171..fc1678aae22b9 100644 --- a/crates/ruff/src/rules/flake8_annotations/rules/definition.rs +++ b/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);