[flake8-async] Update ASYNC109 to match upstream

crates/ruff_linter/resources/test/fixtures/flake8_async/ASYNC109_1.py

@@ -0,0 +1,10 @@
+async def func():
+    ...
+
+
+async def func(timeout):
+    ...
+
+
+async def func(timeout=10):
+    ...

crates/ruff_linter/src/checkers/ast/analyze/statement.rs

@@ -356,7 +356,7 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
                     flake8_builtins::rules::builtin_variable_shadowing(checker, name, name.range());
                 }
             }
-            if checker.enabled(Rule::TrioAsyncFunctionWithTimeout) {
+            if checker.enabled(Rule::AsyncFunctionWithTimeout) {
                 flake8_async::rules::async_function_with_timeout(checker, function_def);
             }
             if checker.enabled(Rule::ReimplementedOperator) {

crates/ruff_linter/src/codes.rs

@@ -295,7 +295,7 @@ pub fn code_to_rule(linter: Linter, code: &str) -> Option<(RuleGroup, Rule)> {
         // flake8-async
         (Flake8Async, "100") => (RuleGroup::Stable, rules::flake8_async::rules::TrioTimeoutWithoutAwait),
         (Flake8Async, "105") => (RuleGroup::Stable, rules::flake8_async::rules::TrioSyncCall),
-        (Flake8Async, "109") => (RuleGroup::Stable, rules::flake8_async::rules::TrioAsyncFunctionWithTimeout),
+        (Flake8Async, "109") => (RuleGroup::Stable, rules::flake8_async::rules::AsyncFunctionWithTimeout),
         (Flake8Async, "110") => (RuleGroup::Stable, rules::flake8_async::rules::TrioUnneededSleep),
         (Flake8Async, "115") => (RuleGroup::Stable, rules::flake8_async::rules::TrioZeroSleepCall),
         (Flake8Async, "116") => (RuleGroup::Preview, rules::flake8_async::rules::SleepForeverCall),

crates/ruff_linter/src/rules/flake8_async/helpers.rs

@@ -1,5 +1,15 @@
 use ruff_python_ast::name::QualifiedName;
 
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+pub(super) enum AsyncModule {
+    /// `anyio`
+    AnyIo,
+    /// `asyncio`
+    AsyncIo,
+    /// `trio`
+    Trio,
+}
+
 #[derive(Debug, Copy, Clone, PartialEq, Eq)]
 pub(super) enum MethodName {
     AcloseForcefully,

crates/ruff_linter/src/rules/flake8_async/mod.rs

@@ -9,14 +9,16 @@ mod tests {
     use anyhow::Result;
     use test_case::test_case;
 
-    use crate::assert_messages;
     use crate::registry::Rule;
+    use crate::settings::types::PreviewMode;
     use crate::settings::LinterSettings;
     use crate::test::test_path;
+    use crate::{assert_messages, settings};
 
     #[test_case(Rule::TrioTimeoutWithoutAwait, Path::new("ASYNC100.py"))]
     #[test_case(Rule::TrioSyncCall, Path::new("ASYNC105.py"))]
-    #[test_case(Rule::TrioAsyncFunctionWithTimeout, Path::new("ASYNC109.py"))]
+    #[test_case(Rule::AsyncFunctionWithTimeout, Path::new("ASYNC109_0.py"))]
+    #[test_case(Rule::AsyncFunctionWithTimeout, Path::new("ASYNC109_1.py"))]
     #[test_case(Rule::TrioUnneededSleep, Path::new("ASYNC110.py"))]
     #[test_case(Rule::TrioZeroSleepCall, Path::new("ASYNC115.py"))]
     #[test_case(Rule::SleepForeverCall, Path::new("ASYNC116.py"))]
@@ -35,4 +37,23 @@ mod tests {
         assert_messages!(snapshot, diagnostics);
         Ok(())
     }
+
+    #[test_case(Rule::AsyncFunctionWithTimeout, Path::new("ASYNC109_0.py"))]
+    #[test_case(Rule::AsyncFunctionWithTimeout, Path::new("ASYNC109_1.py"))]
+    fn preview_rules(rule_code: Rule, path: &Path) -> Result<()> {
+        let snapshot = format!(
+            "preview__{}_{}",
+            rule_code.noqa_code(),
+            path.to_string_lossy()
+        );
+        let diagnostics = test_path(
+            Path::new("flake8_async").join(path).as_path(),
+            &settings::LinterSettings {
+                preview: PreviewMode::Enabled,
+                ..settings::LinterSettings::for_rule(rule_code)
+            },
+        )?;
+        assert_messages!(snapshot, diagnostics);
+        Ok(())
+    }
 }

crates/ruff_linter/src/rules/flake8_async/rules/async_function_with_timeout.rs

@@ -5,38 +5,60 @@ use ruff_python_semantic::Modules;
 use ruff_text_size::Ranged;
 
 use crate::checkers::ast::Checker;
+use crate::rules::flake8_async::helpers::AsyncModule;
+use crate::settings::types::PreviewMode;
 
 /// ## What it does
 /// Checks for `async` functions with a `timeout` argument.
 ///
 /// ## Why is this bad?
 /// Rather than implementing asynchronous timeout behavior manually, prefer
-/// trio's built-in timeout functionality, available as `trio.fail_after`,
-/// `trio.move_on_after`, `trio.fail_at`, and `trio.move_on_at`.
-///
-/// ## Known problems
-/// To avoid false positives, this rule is only enabled if `trio` is imported
-/// in the module.
+/// built-in timeout functionality, such as `asyncio.timeout`, `trio.fail_after`,
+/// or `anyio.move_on_after`, among others.
 ///
 /// ## Example
 /// ```python
-/// async def func():
+/// async def long_running_task(timeout):
+///     ...
+///
+///
+/// async def main():
 ///     await long_running_task(timeout=2)
 /// ```
 ///
 /// Use instead:
 /// ```python
-/// async def func():
-///     with trio.fail_after(2):
+/// async def long_running_task():
+///     ...
+///
+///
+/// async def main():
+///     with asyncio.timeout(2):
 ///         await long_running_task()
 /// ```
+///
+/// [`asyncio` timeouts]: https://docs.python.org/3/library/asyncio-task.html#timeouts
+/// [`anyio` timeouts]: https://anyio.readthedocs.io/en/stable/cancellation.html
+/// [`trio` timeouts]: https://trio.readthedocs.io/en/stable/reference-core.html#cancellation-and-timeouts
 #[violation]
-pub struct TrioAsyncFunctionWithTimeout;
+pub struct AsyncFunctionWithTimeout {
+    module: AsyncModule,
+}
 
-impl Violation for TrioAsyncFunctionWithTimeout {
+impl Violation for AsyncFunctionWithTimeout {
     #[derive_message_formats]
     fn message(&self) -> String {
-        format!("Prefer `trio.fail_after` and `trio.move_on_after` over manual `async` timeout behavior")
+        format!("Async function definition with a `timeout` parameter")
+    }
+
+    fn fix_title(&self) -> Option<String> {
+        let Self { module } = self;
+        let recommendation = match module {
+            AsyncModule::AnyIo => "anyio.fail_after",
+            AsyncModule::Trio => "trio.fail_after",
+            AsyncModule::AsyncIo => "asyncio.timeout",
+        };
+        Some(format!("Use `{recommendation}` instead"))
     }
 }
 
@@ -50,18 +72,31 @@ pub(crate) fn async_function_with_timeout(
         return;
     }
 
-    // If `trio` isn't in scope, avoid raising the diagnostic.
-    if !checker.semantic().seen_module(Modules::TRIO) {
-        return;
-    }
-
     // If the function doesn't have a `timeout` parameter, avoid raising the diagnostic.
     let Some(timeout) = function_def.parameters.find("timeout") else {
         return;
     };
 
-    checker.diagnostics.push(Diagnostic::new(
-        TrioAsyncFunctionWithTimeout,
-        timeout.range(),
-    ));
+    // Get preferred module.
+    let module = if checker.semantic().seen_module(Modules::ANYIO) {
+        AsyncModule::AnyIo
+    } else if checker.semantic().seen_module(Modules::TRIO) {
+        AsyncModule::Trio
+    } else {
+        AsyncModule::AsyncIo
+    };
+
+    if matches!(checker.settings.preview, PreviewMode::Disabled) {
+        if matches!(module, AsyncModule::Trio) {
+            checker.diagnostics.push(Diagnostic::new(
+                AsyncFunctionWithTimeout { module },
+                timeout.range(),
+            ));
+        }
+    } else {
+        checker.diagnostics.push(Diagnostic::new(
+            AsyncFunctionWithTimeout { module },
+            timeout.range(),
+        ));
+    }
 }

crates/ruff_linter/src/rules/flake8_async/snapshots/ruff_linter__rules__flake8_async__tests__ASYNC109_ASYNC109.py.snap → crates/ruff_linter/src/rules/flake8_async/snapshots/ruff_linter__rules__flake8_async__tests__ASYNC109_ASYNC109_0.py.snap

@@ -1,16 +1,18 @@
 ---
 source: crates/ruff_linter/src/rules/flake8_async/mod.rs
 ---
-ASYNC109.py:8:16: ASYNC109 Prefer `trio.fail_after` and `trio.move_on_after` over manual `async` timeout behavior
+ASYNC109_0.py:8:16: ASYNC109 Async function definition with a `timeout` parameter
   |
 8 | async def func(timeout):
   |                ^^^^^^^ ASYNC109
 9 |     ...
   |
+  = help: Use `trio.fail_after` instead
 
-ASYNC109.py:12:16: ASYNC109 Prefer `trio.fail_after` and `trio.move_on_after` over manual `async` timeout behavior
+ASYNC109_0.py:12:16: ASYNC109 Async function definition with a `timeout` parameter
    |
 12 | async def func(timeout=10):
    |                ^^^^^^^^^^ ASYNC109
 13 |     ...
    |
+   = help: Use `trio.fail_after` instead

crates/ruff_linter/src/rules/flake8_async/snapshots/ruff_linter__rules__flake8_async__tests__ASYNC109_ASYNC109_1.py.snap

@@ -0,0 +1,4 @@
+---
+source: crates/ruff_linter/src/rules/flake8_async/mod.rs
+---
+

crates/ruff_linter/src/rules/flake8_async/snapshots/ruff_linter__rules__flake8_async__tests__preview__ASYNC109_ASYNC109_0.py.snap

@@ -0,0 +1,18 @@
+---
+source: crates/ruff_linter/src/rules/flake8_async/mod.rs
+---
+ASYNC109_0.py:8:16: ASYNC109 Async function definition with a `timeout` parameter
+  |
+8 | async def func(timeout):
+  |                ^^^^^^^ ASYNC109
+9 |     ...
+  |
+  = help: Use `trio.fail_after` instead
+
+ASYNC109_0.py:12:16: ASYNC109 Async function definition with a `timeout` parameter
+   |
+12 | async def func(timeout=10):
+   |                ^^^^^^^^^^ ASYNC109
+13 |     ...
+   |
+   = help: Use `trio.fail_after` instead

crates/ruff_linter/src/rules/flake8_async/snapshots/ruff_linter__rules__flake8_async__tests__preview__ASYNC109_ASYNC109_1.py.snap

@@ -0,0 +1,18 @@
+---
+source: crates/ruff_linter/src/rules/flake8_async/mod.rs
+---
+ASYNC109_1.py:5:16: ASYNC109 Async function definition with a `timeout` parameter
+  |
+5 | async def func(timeout):
+  |                ^^^^^^^ ASYNC109
+6 |     ...
+  |
+  = help: Use `asyncio.timeout` instead
+
+ASYNC109_1.py:9:16: ASYNC109 Async function definition with a `timeout` parameter
+   |
+ 9 | async def func(timeout=10):
+   |                ^^^^^^^^^^ ASYNC109
+10 |     ...
+   |
+   = help: Use `asyncio.timeout` instead

crates/ruff_python_semantic/src/model.rs

@@ -1231,6 +1231,7 @@ impl<'a> SemanticModel<'a> {
     pub fn add_module(&mut self, module: &str) {
         match module {
             "_typeshed" => self.seen.insert(Modules::TYPESHED),
+            "anyio" => self.seen.insert(Modules::ANYIO),
             "builtins" => self.seen.insert(Modules::BUILTINS),
             "collections" => self.seen.insert(Modules::COLLECTIONS),
             "contextvars" => self.seen.insert(Modules::CONTEXTVARS),
@@ -1822,6 +1823,7 @@ bitflags! {
         const DATACLASSES = 1 << 17;
         const BUILTINS = 1 << 18;
         const CONTEXTVARS = 1 << 19;
+        const ANYIO = 1 << 20;
     }
 }