Update return type for PT022 autofix

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

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,
                 );

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) {

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