mirror of
https://github.com/astral-sh/ruff.git
synced 2025-10-02 22:54:42 +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():
|
def error():
|
||||||
resource = acquire_resource()
|
resource = acquire_resource()
|
||||||
yield 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,
|
stmt,
|
||||||
name,
|
name,
|
||||||
parameters,
|
parameters,
|
||||||
|
returns.as_deref(),
|
||||||
decorator_list,
|
decorator_list,
|
||||||
body,
|
body,
|
||||||
);
|
);
|
||||||
|
|
|
@ -764,7 +764,13 @@ fn check_fixture_decorator(checker: &mut Checker, func_name: &str, decorator: &D
|
||||||
}
|
}
|
||||||
|
|
||||||
/// PT004, PT005, PT022
|
/// 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();
|
let mut visitor = SkipFunctionsVisitor::default();
|
||||||
|
|
||||||
for stmt in body {
|
for stmt in body {
|
||||||
|
@ -795,10 +801,18 @@ fn check_fixture_returns(checker: &mut Checker, stmt: &Stmt, name: &str, body: &
|
||||||
}
|
}
|
||||||
|
|
||||||
if checker.enabled(Rule::PytestUselessYieldFixture) {
|
if checker.enabled(Rule::PytestUselessYieldFixture) {
|
||||||
if let Some(stmt) = body.last() {
|
let Some(stmt) = body.last() else {
|
||||||
if let Stmt::Expr(ast::StmtExpr { value, range: _ }) = stmt {
|
return;
|
||||||
if value.is_yield_expr() {
|
};
|
||||||
if visitor.yield_statements.len() == 1 {
|
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(
|
let mut diagnostic = Diagnostic::new(
|
||||||
PytestUselessYieldFixture {
|
PytestUselessYieldFixture {
|
||||||
name: name.to_string(),
|
name: name.to_string(),
|
||||||
|
@ -806,17 +820,32 @@ fn check_fixture_returns(checker: &mut Checker, stmt: &Stmt, name: &str, body: &
|
||||||
stmt.range(),
|
stmt.range(),
|
||||||
);
|
);
|
||||||
if checker.patch(diagnostic.kind.rule()) {
|
if checker.patch(diagnostic.kind.rule()) {
|
||||||
diagnostic.set_fix(Fix::automatic(Edit::range_replacement(
|
let yield_edit = Edit::range_replacement(
|
||||||
"return".to_string(),
|
"return".to_string(),
|
||||||
TextRange::at(stmt.start(), "yield".text_len()),
|
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);
|
checker.diagnostics.push(diagnostic);
|
||||||
}
|
}
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// PT019
|
/// PT019
|
||||||
|
@ -910,6 +939,7 @@ pub(crate) fn fixture(
|
||||||
stmt: &Stmt,
|
stmt: &Stmt,
|
||||||
name: &str,
|
name: &str,
|
||||||
parameters: &Parameters,
|
parameters: &Parameters,
|
||||||
|
returns: Option<&Expr>,
|
||||||
decorators: &[Decorator],
|
decorators: &[Decorator],
|
||||||
body: &[Stmt],
|
body: &[Stmt],
|
||||||
) {
|
) {
|
||||||
|
@ -933,7 +963,7 @@ pub(crate) fn fixture(
|
||||||
|| checker.enabled(Rule::PytestUselessYieldFixture))
|
|| checker.enabled(Rule::PytestUselessYieldFixture))
|
||||||
&& !is_abstract(decorators, checker.semantic())
|
&& !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) {
|
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()
|
16 16 | resource = acquire_resource()
|
||||||
17 |- yield resource
|
17 |- yield resource
|
||||||
17 |+ return 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