From f0ec2303f99ae2e42878517596884565172aa278 Mon Sep 17 00:00:00 2001
From: Micha Reiser <micha@reiser.io>
Date: Wed, 10 Jul 2024 09:22:54 +0200
Subject: [PATCH] Remove `unwrap`, share more code between preview and
 non-preivew

---
 .../flake8_async/rules/async_zero_sleep.rs    | 63 ++++++++-----------
 1 file changed, 25 insertions(+), 38 deletions(-)

diff --git a/crates/ruff_linter/src/rules/flake8_async/rules/async_zero_sleep.rs b/crates/ruff_linter/src/rules/flake8_async/rules/async_zero_sleep.rs
index 906ea5fe0042e..9f9ef57cbd7d6 100644
--- a/crates/ruff_linter/src/rules/flake8_async/rules/async_zero_sleep.rs
+++ b/crates/ruff_linter/src/rules/flake8_async/rules/async_zero_sleep.rs
@@ -7,7 +7,6 @@ use ruff_text_size::Ranged;
 use crate::checkers::ast::Checker;
 use crate::importer::ImportRequest;
 use crate::rules::flake8_async::helpers::AsyncModule;
-use crate::settings::types::PreviewMode;
 
 /// ## What it does
 /// Checks for uses of `trio.sleep(0)` or `anyio.sleep(0)`.
@@ -83,43 +82,31 @@ pub(crate) fn async_zero_sleep(checker: &mut Checker, call: &ExprCall) {
         return;
     };
 
-    if matches!(checker.settings.preview, PreviewMode::Disabled) {
-        if matches!(qualified_name.segments(), ["trio", "sleep"]) {
-            let mut diagnostic = Diagnostic::new(
-                AsyncZeroSleep {
-                    module: AsyncModule::Trio,
-                },
-                call.range(),
-            );
-            diagnostic.try_set_fix(|| {
-                let (import_edit, binding) = checker.importer().get_or_import_symbol(
-                    &ImportRequest::import_from("trio", "lowlevel"),
-                    call.func.start(),
-                    checker.semantic(),
-                )?;
-                let reference_edit =
-                    Edit::range_replacement(format!("{binding}.checkpoint"), call.func.range());
-                let arg_edit = Edit::range_replacement("()".to_string(), call.arguments.range());
-                Ok(Fix::safe_edits(import_edit, [reference_edit, arg_edit]))
-            });
-            checker.diagnostics.push(diagnostic);
-        }
-    } else {
-        if matches!(qualified_name.segments(), ["trio" | "anyio", "sleep"]) {
-            let module = AsyncModule::try_from(&qualified_name).unwrap();
-            let mut diagnostic = Diagnostic::new(AsyncZeroSleep { module }, call.range());
-            diagnostic.try_set_fix(|| {
-                let (import_edit, binding) = checker.importer().get_or_import_symbol(
-                    &ImportRequest::import_from(&module.to_string(), "lowlevel"),
-                    call.func.start(),
-                    checker.semantic(),
-                )?;
-                let reference_edit =
-                    Edit::range_replacement(format!("{binding}.checkpoint"), call.func.range());
-                let arg_edit = Edit::range_replacement("()".to_string(), call.arguments.range());
-                Ok(Fix::safe_edits(import_edit, [reference_edit, arg_edit]))
-            });
-            checker.diagnostics.push(diagnostic);
+    if let Some(module) = AsyncModule::try_from(&qualified_name) {
+        let is_relevant_module = if checker.settings.preview.is_enabled() {
+            matches!(module, AsyncModule::Trio | AsyncModule::AnyIo)
+        } else {
+            matches!(module, AsyncModule::Trio)
+        };
+
+        let is_sleep = is_relevant_module && matches!(qualified_name.segments(), [_, "sleep"]);
+
+        if !is_sleep {
+            return;
         }
+
+        let mut diagnostic = Diagnostic::new(AsyncZeroSleep { module }, call.range());
+        diagnostic.try_set_fix(|| {
+            let (import_edit, binding) = checker.importer().get_or_import_symbol(
+                &ImportRequest::import_from(&module.to_string(), "lowlevel"),
+                call.func.start(),
+                checker.semantic(),
+            )?;
+            let reference_edit =
+                Edit::range_replacement(format!("{binding}.checkpoint"), call.func.range());
+            let arg_edit = Edit::range_replacement("()".to_string(), call.arguments.range());
+            Ok(Fix::safe_edits(import_edit, [reference_edit, arg_edit]))
+        });
+        checker.diagnostics.push(diagnostic);
     }
 }