From 46f33b4fbb2d0af67d60998d1dfa031229e3778f Mon Sep 17 00:00:00 2001 From: Dhruv Manilawala Date: Sat, 23 Sep 2023 10:48:29 +0530 Subject: [PATCH] Update return type for `PT022` autofix --- .../fixtures/flake8_pytest_style/PT022.py | 26 +++++++ .../src/checkers/ast/analyze/statement.rs | 1 + .../flake8_pytest_style/rules/fixture.rs | 70 +++++++++++++------ ...es__flake8_pytest_style__tests__PT022.snap | 44 ++++++++++++ 4 files changed, 121 insertions(+), 20 deletions(-) diff --git a/crates/ruff_linter/resources/test/fixtures/flake8_pytest_style/PT022.py b/crates/ruff_linter/resources/test/fixtures/flake8_pytest_style/PT022.py index 89c64da013c80..dc059585193c2 100644 --- a/crates/ruff_linter/resources/test/fixtures/flake8_pytest_style/PT022.py +++ b/crates/ruff_linter/resources/test/fixtures/flake8_pytest_style/PT022.py @@ -15,3 +15,29 @@ def ok_complex_logic(): def error(): resource = acquire_resource() yield resource + + +import typing +from typing import Generator + + +@pytest.fixture() +def ok_complex_logic() -> typing.Generator[Resource, None, None]: + if some_condition: + resource = acquire_resource() + yield resource + resource.release() + return + yield None + + +@pytest.fixture() +def error() -> typing.Generator[typing.Any, None, None]: + resource = acquire_resource() + yield resource + + +@pytest.fixture() +def error() -> Generator[Resource, None, None]: + resource = acquire_resource() + yield resource diff --git a/crates/ruff_linter/src/checkers/ast/analyze/statement.rs b/crates/ruff_linter/src/checkers/ast/analyze/statement.rs index 946d4f5c08dc1..a39fd0734355c 100644 --- a/crates/ruff_linter/src/checkers/ast/analyze/statement.rs +++ b/crates/ruff_linter/src/checkers/ast/analyze/statement.rs @@ -292,6 +292,7 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) { stmt, name, parameters, + returns.as_deref(), decorator_list, body, ); diff --git a/crates/ruff_linter/src/rules/flake8_pytest_style/rules/fixture.rs b/crates/ruff_linter/src/rules/flake8_pytest_style/rules/fixture.rs index 4e401c6b44d12..b3374ba07ee5d 100644 --- a/crates/ruff_linter/src/rules/flake8_pytest_style/rules/fixture.rs +++ b/crates/ruff_linter/src/rules/flake8_pytest_style/rules/fixture.rs @@ -764,7 +764,13 @@ fn check_fixture_decorator(checker: &mut Checker, func_name: &str, decorator: &D } /// PT004, PT005, PT022 -fn check_fixture_returns(checker: &mut Checker, stmt: &Stmt, name: &str, body: &[Stmt]) { +fn check_fixture_returns( + checker: &mut Checker, + stmt: &Stmt, + name: &str, + body: &[Stmt], + returns: Option<&Expr>, +) { let mut visitor = SkipFunctionsVisitor::default(); for stmt in body { @@ -795,27 +801,50 @@ fn check_fixture_returns(checker: &mut Checker, stmt: &Stmt, name: &str, body: & } if checker.enabled(Rule::PytestUselessYieldFixture) { - if let Some(stmt) = body.last() { - if let Stmt::Expr(ast::StmtExpr { value, range: _ }) = stmt { - if value.is_yield_expr() { - if visitor.yield_statements.len() == 1 { - let mut diagnostic = Diagnostic::new( - PytestUselessYieldFixture { - name: name.to_string(), - }, - stmt.range(), - ); - if checker.patch(diagnostic.kind.rule()) { - diagnostic.set_fix(Fix::automatic(Edit::range_replacement( - "return".to_string(), - TextRange::at(stmt.start(), "yield".text_len()), - ))); - } - checker.diagnostics.push(diagnostic); - } + let Some(stmt) = body.last() else { + return; + }; + let Stmt::Expr(ast::StmtExpr { value, .. }) = stmt else { + return; + }; + if !value.is_yield_expr() { + return; + } + if visitor.yield_statements.len() != 1 { + return; + } + let mut diagnostic = Diagnostic::new( + PytestUselessYieldFixture { + name: name.to_string(), + }, + stmt.range(), + ); + if checker.patch(diagnostic.kind.rule()) { + let yield_edit = Edit::range_replacement( + "return".to_string(), + TextRange::at(stmt.start(), "yield".text_len()), + ); + let return_type_edit = returns.and_then(|returns| { + let ast::ExprSubscript { value, slice, .. } = returns.as_subscript_expr()?; + let ast::ExprTuple { elts, .. } = slice.as_tuple_expr()?; + let [first, ..] = elts.as_slice() else { + return None; + }; + if !checker.semantic().match_typing_expr(value, "Generator") { + return None; } + Some(Edit::range_replacement( + checker.generator().expr(first), + returns.range(), + )) + }); + if let Some(return_type_edit) = return_type_edit { + diagnostic.set_fix(Fix::automatic_edits(yield_edit, [return_type_edit])); + } else { + diagnostic.set_fix(Fix::automatic(yield_edit)); } } + checker.diagnostics.push(diagnostic); } } @@ -910,6 +939,7 @@ pub(crate) fn fixture( stmt: &Stmt, name: &str, parameters: &Parameters, + returns: Option<&Expr>, decorators: &[Decorator], body: &[Stmt], ) { @@ -933,7 +963,7 @@ pub(crate) fn fixture( || checker.enabled(Rule::PytestUselessYieldFixture)) && !is_abstract(decorators, checker.semantic()) { - check_fixture_returns(checker, stmt, name, body); + check_fixture_returns(checker, stmt, name, body, returns); } if checker.enabled(Rule::PytestFixtureFinalizerCallback) { diff --git a/crates/ruff_linter/src/rules/flake8_pytest_style/snapshots/ruff_linter__rules__flake8_pytest_style__tests__PT022.snap b/crates/ruff_linter/src/rules/flake8_pytest_style/snapshots/ruff_linter__rules__flake8_pytest_style__tests__PT022.snap index ea497c7f5aacf..b9134a42d54bf 100644 --- a/crates/ruff_linter/src/rules/flake8_pytest_style/snapshots/ruff_linter__rules__flake8_pytest_style__tests__PT022.snap +++ b/crates/ruff_linter/src/rules/flake8_pytest_style/snapshots/ruff_linter__rules__flake8_pytest_style__tests__PT022.snap @@ -16,5 +16,49 @@ PT022.py:17:5: PT022 [*] No teardown in fixture `error`, use `return` instead of 16 16 | resource = acquire_resource() 17 |- yield resource 17 |+ return resource +18 18 | +19 19 | +20 20 | import typing + +PT022.py:37:5: PT022 [*] No teardown in fixture `error`, use `return` instead of `yield` + | +35 | def error() -> typing.Generator[typing.Any, None, None]: +36 | resource = acquire_resource() +37 | yield resource + | ^^^^^^^^^^^^^^ PT022 + | + = help: Replace `yield` with `return` + +ℹ Fix +32 32 | +33 33 | +34 34 | @pytest.fixture() +35 |-def error() -> typing.Generator[typing.Any, None, None]: + 35 |+def error() -> typing.Any: +36 36 | resource = acquire_resource() +37 |- yield resource + 37 |+ return resource +38 38 | +39 39 | +40 40 | @pytest.fixture() + +PT022.py:43:5: PT022 [*] No teardown in fixture `error`, use `return` instead of `yield` + | +41 | def error() -> Generator[Resource, None, None]: +42 | resource = acquire_resource() +43 | yield resource + | ^^^^^^^^^^^^^^ PT022 + | + = help: Replace `yield` with `return` + +ℹ Fix +38 38 | +39 39 | +40 40 | @pytest.fixture() +41 |-def error() -> Generator[Resource, None, None]: + 41 |+def error() -> Resource: +42 42 | resource = acquire_resource() +43 |- yield resource + 43 |+ return resource