From bcdb3f984091be3c566908f8c01d2a36c353eb7e Mon Sep 17 00:00:00 2001
From: Alex Waygood <Alex.Waygood@Gmail.com>
Date: Sat, 1 Feb 2025 15:44:42 +0000
Subject: [PATCH] Use `Diagnostic::try_set_fix` in `bad-generator-return-type`
 (#15873)

---
 .../rules/bad_generator_return_type.rs        | 39 +++++++++----------
 1 file changed, 19 insertions(+), 20 deletions(-)

diff --git a/crates/ruff_linter/src/rules/flake8_pyi/rules/bad_generator_return_type.rs b/crates/ruff_linter/src/rules/flake8_pyi/rules/bad_generator_return_type.rs
index 7aad7fb479..c4f97c1bf4 100644
--- a/crates/ruff_linter/src/rules/flake8_pyi/rules/bad_generator_return_type.rs
+++ b/crates/ruff_linter/src/rules/flake8_pyi/rules/bad_generator_return_type.rs
@@ -220,16 +220,18 @@ pub(crate) fn bad_generator_return_type(
         },
         function_def.identifier(),
     );
-    if let Some(fix) = generate_fix(
-        function_def,
-        returns,
-        yield_type_info,
-        module,
-        member,
-        checker,
-    ) {
-        diagnostic.set_fix(fix);
-    };
+
+    diagnostic.try_set_fix(|| {
+        generate_fix(
+            function_def,
+            returns,
+            yield_type_info,
+            module,
+            member,
+            checker,
+        )
+    });
+
     checker.diagnostics.push(diagnostic);
 }
 
@@ -246,17 +248,14 @@ fn generate_fix(
     module: Module,
     member: Generator,
     checker: &Checker,
-) -> Option<Fix> {
+) -> anyhow::Result<Fix> {
     let expr = map_subscript(returns);
 
-    let (import_edit, binding) = checker
-        .importer()
-        .get_or_import_symbol(
-            &ImportRequest::import_from(&module.to_string(), &member.to_iter().to_string()),
-            expr.start(),
-            checker.semantic(),
-        )
-        .ok()?;
+    let (import_edit, binding) = checker.importer().get_or_import_symbol(
+        &ImportRequest::import_from(&module.to_string(), &member.to_iter().to_string()),
+        expr.start(),
+        checker.semantic(),
+    )?;
     let binding_edit = Edit::range_replacement(binding, expr.range());
     let yield_edit = yield_type_info.map(|yield_type_info| {
         Edit::range_replacement(
@@ -272,7 +271,7 @@ fn generate_fix(
         Applicability::Unsafe
     };
 
-    Some(Fix::applicable_edits(
+    Ok(Fix::applicable_edits(
         import_edit,
         std::iter::once(binding_edit).chain(yield_edit),
         applicability,