Skip to content

Commit

Permalink
Update return type for PT022 autofix (#7613)
Browse files Browse the repository at this point in the history
## Summary

This PR fixes the autofix behavior for `PT022` to create an additional
edit for the return type if it's present. The edit will update the
return type from `Generator[T, ...]` to `T`. As per the [official
documentation](https://docs.python.org/3/library/typing.html?highlight=typing%20generator#typing.Generator),
the first position is the yield type, so we can ignore other positions.

```python
typing.Generator[YieldType, SendType, ReturnType]
```

## Test Plan

Add new test cases, `cargo test` and review the snapshots.

fixes: #7610
  • Loading branch information
dhruvmanila authored Sep 24, 2023
1 parent 604cf52 commit 15813a6
Show file tree
Hide file tree
Showing 4 changed files with 121 additions and 20 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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
1 change: 1 addition & 0 deletions crates/ruff_linter/src/checkers/ast/analyze/statement.rs
Original file line number Diff line number Diff line change
Expand Up @@ -292,6 +292,7 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
stmt,
name,
parameters,
returns.as_deref(),
decorator_list,
body,
);
Expand Down
70 changes: 50 additions & 20 deletions crates/ruff_linter/src/rules/flake8_pytest_style/rules/fixture.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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);
}
}

Expand Down Expand Up @@ -910,6 +939,7 @@ pub(crate) fn fixture(
stmt: &Stmt,
name: &str,
parameters: &Parameters,
returns: Option<&Expr>,
decorators: &[Decorator],
body: &[Stmt],
) {
Expand All @@ -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) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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


0 comments on commit 15813a6

Please sign in to comment.