mirror of
https://github.com/astral-sh/ruff.git
synced 2025-10-01 22:31:23 +00:00
Update return type for PT022
autofix (#7613)
## 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
This commit is contained in:
parent
604cf521b5
commit
15813a65f3
4 changed files with 121 additions and 20 deletions
|
@ -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
|
||||
|
|
|
@ -292,6 +292,7 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
|
|||
stmt,
|
||||
name,
|
||||
parameters,
|
||||
returns.as_deref(),
|
||||
decorator_list,
|
||||
body,
|
||||
);
|
||||
|
|
|
@ -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) {
|
||||
|
|
|
@ -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
|
||||
|
||||
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue