[flake8_async] Refactor argument name resolution for async sleep func… (#18262)

Co-authored-by: Vasanth-96 <ramavath.naik@itilite.com>
Co-authored-by: Micha Reiser <micha@reiser.io>

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

@@ -145,3 +145,23 @@ def func():
             sleep = 10
 
     anyio.run(main)
+
+
+async def test_anyio_async115_helpers():
+    import anyio
+
+    await anyio.sleep(delay=1)  # OK
+    await anyio.sleep(seconds=1)  # OK
+
+    await anyio.sleep(delay=0)  # ASYNC115
+    await anyio.sleep(seconds=0)  # OK
+
+
+async def test_trio_async115_helpers():
+    import trio
+
+    await trio.sleep(seconds=1)  # OK
+    await trio.sleep(delay=1)  # OK
+
+    await trio.sleep(seconds=0)  # ASYNC115
+    await trio.sleep(delay=0)  # OK

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

@@ -108,3 +108,23 @@ async def import_from_anyio():
 
     # catch from import
     await sleep(86401)  # error: 116, "async"
+
+
+async def test_anyio_async116_helpers():
+    import anyio
+
+    await anyio.sleep(delay=1)  # OK
+    await anyio.sleep(seconds=1)  # OK
+
+    await anyio.sleep(delay=86401)  # ASYNC116
+    await anyio.sleep(seconds=86401)  # OK
+
+
+async def test_trio_async116_helpers():
+    import trio
+
+    await trio.sleep(seconds=1)  # OK
+    await trio.sleep(delay=1)  # OK
+
+    await trio.sleep(seconds=86401)  # ASYNC116
+    await trio.sleep(delay=86401)  # OK

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

@@ -62,7 +62,25 @@ pub(crate) fn async_zero_sleep(checker: &Checker, call: &ExprCall) {
         return;
     }
 
-    let Some(arg) = call.arguments.find_argument_value("seconds", 0) else {
+    let Some(qualified_name) = checker
+        .semantic()
+        .resolve_qualified_name(call.func.as_ref())
+    else {
+        return;
+    };
+
+    let Some(module) = AsyncModule::try_from(&qualified_name) else {
+        return;
+    };
+
+    // Determine the correct argument name
+    let arg_name = match module {
+        AsyncModule::Trio => "seconds",
+        AsyncModule::AnyIo => "delay",
+        AsyncModule::AsyncIo => return,
+    };
+
+    let Some(arg) = call.arguments.find_argument_value(arg_name, 0) else {
         return;
     };
 

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

@@ -71,7 +71,25 @@ pub(crate) fn long_sleep_not_forever(checker: &Checker, call: &ExprCall) {
         return;
     }
 
-    let Some(arg) = call.arguments.find_argument_value("seconds", 0) else {
+    let Some(qualified_name) = checker
+        .semantic()
+        .resolve_qualified_name(call.func.as_ref())
+    else {
+        return;
+    };
+
+    let Some(module) = AsyncModule::try_from(&qualified_name) else {
+        return;
+    };
+
+    // Determine the correct argument name
+    let arg_name = match module {
+        AsyncModule::Trio => "seconds",
+        AsyncModule::AnyIo => "delay",
+        AsyncModule::AsyncIo => return,
+    };
+
+    let Some(arg) = call.arguments.find_argument_value(arg_name, 0) else {
         return;
     };
 

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

@@ -214,3 +214,41 @@ ASYNC115.py:128:15: ASYNC115 [*] Use `anyio.lowlevel.checkpoint()` instead of `a
 129 129 | 
 130 130 | 
 131 131 | def func():
+
+ASYNC115.py:156:11: ASYNC115 [*] Use `anyio.lowlevel.checkpoint()` instead of `anyio.sleep(0)`
+    |
+154 |     await anyio.sleep(seconds=1)  # OK
+155 |
+156 |     await anyio.sleep(delay=0)  # ASYNC115
+    |           ^^^^^^^^^^^^^^^^^^^^ ASYNC115
+157 |     await anyio.sleep(seconds=0)  # OK
+    |
+    = help: Replace with `anyio.lowlevel.checkpoint()`
+
+ℹ Safe fix
+153 153 |     await anyio.sleep(delay=1)  # OK
+154 154 |     await anyio.sleep(seconds=1)  # OK
+155 155 | 
+156     |-    await anyio.sleep(delay=0)  # ASYNC115
+    156 |+    await anyio.lowlevel.checkpoint()  # ASYNC115
+157 157 |     await anyio.sleep(seconds=0)  # OK
+158 158 | 
+159 159 | 
+
+ASYNC115.py:166:11: ASYNC115 [*] Use `trio.lowlevel.checkpoint()` instead of `trio.sleep(0)`
+    |
+164 |     await trio.sleep(delay=1)  # OK
+165 |
+166 |     await trio.sleep(seconds=0)  # ASYNC115
+    |           ^^^^^^^^^^^^^^^^^^^^^ ASYNC115
+167 |     await trio.sleep(delay=0)  # OK
+    |
+    = help: Replace with `trio.lowlevel.checkpoint()`
+
+ℹ Safe fix
+163 163 |     await trio.sleep(seconds=1)  # OK
+164 164 |     await trio.sleep(delay=1)  # OK
+165 165 | 
+166     |-    await trio.sleep(seconds=0)  # ASYNC115
+    166 |+    await trio.lowlevel.checkpoint()  # ASYNC115
+167 167 |     await trio.sleep(delay=0)  # OK

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

@@ -289,3 +289,44 @@ ASYNC116.py:110:11: ASYNC116 [*] `anyio.sleep()` with >24 hour interval should u
 109 110 |     # catch from import
 110     |-    await sleep(86401)  # error: 116, "async"
     111 |+    await sleep_forever()  # error: 116, "async"
+111 112 | 
+112 113 | 
+113 114 | async def test_anyio_async116_helpers():
+
+ASYNC116.py:119:11: ASYNC116 [*] `anyio.sleep()` with >24 hour interval should usually be `anyio.sleep_forever()`
+    |
+117 |     await anyio.sleep(seconds=1)  # OK
+118 |
+119 |     await anyio.sleep(delay=86401)  # ASYNC116
+    |           ^^^^^^^^^^^^^^^^^^^^^^^^ ASYNC116
+120 |     await anyio.sleep(seconds=86401)  # OK
+    |
+    = help: Replace with `anyio.sleep_forever()`
+
+ℹ Unsafe fix
+116 116 |     await anyio.sleep(delay=1)  # OK
+117 117 |     await anyio.sleep(seconds=1)  # OK
+118 118 | 
+119     |-    await anyio.sleep(delay=86401)  # ASYNC116
+    119 |+    await anyio.sleep_forever()  # ASYNC116
+120 120 |     await anyio.sleep(seconds=86401)  # OK
+121 121 | 
+122 122 | 
+
+ASYNC116.py:129:11: ASYNC116 [*] `trio.sleep()` with >24 hour interval should usually be `trio.sleep_forever()`
+    |
+127 |     await trio.sleep(delay=1)  # OK
+128 |
+129 |     await trio.sleep(seconds=86401)  # ASYNC116
+    |           ^^^^^^^^^^^^^^^^^^^^^^^^^ ASYNC116
+130 |     await trio.sleep(delay=86401)  # OK
+    |
+    = help: Replace with `trio.sleep_forever()`
+
+ℹ Unsafe fix
+126 126 |     await trio.sleep(seconds=1)  # OK
+127 127 |     await trio.sleep(delay=1)  # OK
+128 128 | 
+129     |-    await trio.sleep(seconds=86401)  # ASYNC116
+    129 |+    await trio.sleep_forever()  # ASYNC116
+130 130 |     await trio.sleep(delay=86401)  # OK