Avoid syntax error via importing trio.lowlevel (#8730)

We ended up with a syntax error here via `from trio import lowlevel.checkpoint`. The new solution avoids that error, but does miss cases like: ```py from trio.lowlevel import Timer ``` Where it could insert `from trio.lowlevel import Timer, checkpoint`. Instead, it'll add `from trio import lowlevel`. See: #8402 (comment)
astral-sh · Nov 17, 2023 · 1fcccf8 · 1fcccf8
1 parent 14e65af
commit 1fcccf8
Show file tree

Hide file tree

Showing 3 changed files with 77 additions and 86 deletions.
diff --git a/crates/ruff_linter/resources/test/fixtures/flake8_trio/TRIO115.py b/crates/ruff_linter/resources/test/fixtures/flake8_trio/TRIO115.py
@@ -1,8 +1,7 @@
-import trio
-from trio import sleep
-
-
 async def func():
+    import trio
+    from trio import sleep
+
     await trio.sleep(0)  # TRIO115
     await trio.sleep(1)  # OK
     await trio.sleep(0, 1)  # OK
@@ -21,8 +20,11 @@ async def func():
     trio.sleep(bar)
 
 
-trio.sleep(0)  # TRIO115
+def func():
+    trio.run(trio.sleep(0))  # TRIO115
+
 
+from trio import Event, sleep
 
 def func():
-    trio.run(trio.sleep(0))  # TRIO115
+    sleep(0)  # TRIO115
diff --git a/crates/ruff_linter/src/rules/flake8_trio/rules/zero_sleep_call.rs b/crates/ruff_linter/src/rules/flake8_trio/rules/zero_sleep_call.rs
@@ -97,11 +97,12 @@ pub(crate) fn zero_sleep_call(checker: &mut Checker, call: &ExprCall) {
     let mut diagnostic = Diagnostic::new(TrioZeroSleepCall, call.range());
     diagnostic.try_set_fix(|| {
         let (import_edit, binding) = checker.importer().get_or_import_symbol(
-            &ImportRequest::import("trio", "lowlevel.checkpoint"),
+            &ImportRequest::import_from("trio", "lowlevel"),
             call.func.start(),
             checker.semantic(),
         )?;
-        let reference_edit = Edit::range_replacement(binding, call.func.range());
+        let reference_edit =
+            Edit::range_replacement(format!("{binding}.checkpoint"), call.func.range());
         let arg_edit = Edit::range_deletion(call.arguments.range);
         Ok(Fix::safe_edits(import_edit, [reference_edit, arg_edit]))
     });

diff --git a/...les/flake8_trio/snapshots/ruff_linter__rules__flake8_trio__tests__TRIO115_TRIO115.py.snap b/...les/flake8_trio/snapshots/ruff_linter__rules__flake8_trio__tests__TRIO115_TRIO115.py.snap
@@ -1,119 +1,107 @@
 ---
 source: crates/ruff_linter/src/rules/flake8_trio/mod.rs
 ---
-TRIO115.py:6:11: TRIO115 [*] Use `trio.lowlevel.checkpoint()` instead of `trio.sleep(0)`
+TRIO115.py:5:11: TRIO115 [*] Use `trio.lowlevel.checkpoint()` instead of `trio.sleep(0)`
   |
-5 | async def func():
-6 |     await trio.sleep(0)  # TRIO115
+3 |     from trio import sleep
+4 | 
+5 |     await trio.sleep(0)  # TRIO115
   |           ^^^^^^^^^^^^^ TRIO115
-7 |     await trio.sleep(1)  # OK
-8 |     await trio.sleep(0, 1)  # OK
+6 |     await trio.sleep(1)  # OK
+7 |     await trio.sleep(0, 1)  # OK
   |
   = help: Replace with `trio.lowlevel.checkpoint()`
 
 ℹ Safe fix
-3 3 | 
+2 2 |     import trio
+3 3 |     from trio import sleep
 4 4 | 
-5 5 | async def func():
-6   |-    await trio.sleep(0)  # TRIO115
-  6 |+    await trio.lowlevel.checkpoint  # TRIO115
-7 7 |     await trio.sleep(1)  # OK
-8 8 |     await trio.sleep(0, 1)  # OK
-9 9 |     await trio.sleep(...)  # OK
+5   |-    await trio.sleep(0)  # TRIO115
+  5 |+    await trio.lowlevel.checkpoint  # TRIO115
+6 6 |     await trio.sleep(1)  # OK
+7 7 |     await trio.sleep(0, 1)  # OK
+8 8 |     await trio.sleep(...)  # OK
 
-TRIO115.py:12:5: TRIO115 [*] Use `trio.lowlevel.checkpoint()` instead of `trio.sleep(0)`
+TRIO115.py:11:5: TRIO115 [*] Use `trio.lowlevel.checkpoint()` instead of `trio.sleep(0)`
    |
-10 |     await trio.sleep()  # OK
-11 | 
-12 |     trio.sleep(0)  # TRIO115
+ 9 |     await trio.sleep()  # OK
+10 | 
+11 |     trio.sleep(0)  # TRIO115
    |     ^^^^^^^^^^^^^ TRIO115
-13 |     foo = 0
-14 |     trio.sleep(foo)  # TRIO115
+12 |     foo = 0
+13 |     trio.sleep(foo)  # TRIO115
    |
    = help: Replace with `trio.lowlevel.checkpoint()`
 
 ℹ Safe fix
-9  9  |     await trio.sleep(...)  # OK
-10 10 |     await trio.sleep()  # OK
-11 11 | 
-12    |-    trio.sleep(0)  # TRIO115
-   12 |+    trio.lowlevel.checkpoint  # TRIO115
-13 13 |     foo = 0
-14 14 |     trio.sleep(foo)  # TRIO115
-15 15 |     trio.sleep(1)  # OK
+8  8  |     await trio.sleep(...)  # OK
+9  9  |     await trio.sleep()  # OK
+10 10 | 
+11    |-    trio.sleep(0)  # TRIO115
+   11 |+    trio.lowlevel.checkpoint  # TRIO115
+12 12 |     foo = 0
+13 13 |     trio.sleep(foo)  # TRIO115
+14 14 |     trio.sleep(1)  # OK
 
-TRIO115.py:14:5: TRIO115 [*] Use `trio.lowlevel.checkpoint()` instead of `trio.sleep(0)`
+TRIO115.py:13:5: TRIO115 [*] Use `trio.lowlevel.checkpoint()` instead of `trio.sleep(0)`
    |
-12 |     trio.sleep(0)  # TRIO115
-13 |     foo = 0
-14 |     trio.sleep(foo)  # TRIO115
+11 |     trio.sleep(0)  # TRIO115
+12 |     foo = 0
+13 |     trio.sleep(foo)  # TRIO115
    |     ^^^^^^^^^^^^^^^ TRIO115
-15 |     trio.sleep(1)  # OK
-16 |     time.sleep(0)  # OK
+14 |     trio.sleep(1)  # OK
+15 |     time.sleep(0)  # OK
    |
    = help: Replace with `trio.lowlevel.checkpoint()`
 
 ℹ Safe fix
-11 11 | 
-12 12 |     trio.sleep(0)  # TRIO115
-13 13 |     foo = 0
-14    |-    trio.sleep(foo)  # TRIO115
-   14 |+    trio.lowlevel.checkpoint  # TRIO115
-15 15 |     trio.sleep(1)  # OK
-16 16 |     time.sleep(0)  # OK
-17 17 | 
+10 10 | 
+11 11 |     trio.sleep(0)  # TRIO115
+12 12 |     foo = 0
+13    |-    trio.sleep(foo)  # TRIO115
+   13 |+    trio.lowlevel.checkpoint  # TRIO115
+14 14 |     trio.sleep(1)  # OK
+15 15 |     time.sleep(0)  # OK
+16 16 | 
 
-TRIO115.py:18:5: TRIO115 [*] Use `trio.lowlevel.checkpoint()` instead of `trio.sleep(0)`
+TRIO115.py:17:5: TRIO115 [*] Use `trio.lowlevel.checkpoint()` instead of `trio.sleep(0)`
    |
-16 |     time.sleep(0)  # OK
-17 | 
-18 |     sleep(0)  # TRIO115
+15 |     time.sleep(0)  # OK
+16 | 
+17 |     sleep(0)  # TRIO115
    |     ^^^^^^^^ TRIO115
-19 | 
-20 |     bar = "bar"
+18 | 
+19 |     bar = "bar"
    |
    = help: Replace with `trio.lowlevel.checkpoint()`
 
 ℹ Safe fix
-15 15 |     trio.sleep(1)  # OK
-16 16 |     time.sleep(0)  # OK
-17 17 | 
-18    |-    sleep(0)  # TRIO115
-   18 |+    trio.lowlevel.checkpoint  # TRIO115
-19 19 | 
-20 20 |     bar = "bar"
-21 21 |     trio.sleep(bar)
+14 14 |     trio.sleep(1)  # OK
+15 15 |     time.sleep(0)  # OK
+16 16 | 
+17    |-    sleep(0)  # TRIO115
+   17 |+    trio.lowlevel.checkpoint  # TRIO115
+18 18 | 
+19 19 |     bar = "bar"
+20 20 |     trio.sleep(bar)
 
-TRIO115.py:24:1: TRIO115 [*] Use `trio.lowlevel.checkpoint()` instead of `trio.sleep(0)`
+TRIO115.py:30:5: TRIO115 [*] Use `trio.lowlevel.checkpoint()` instead of `trio.sleep(0)`
    |
-24 | trio.sleep(0)  # TRIO115
-   | ^^^^^^^^^^^^^ TRIO115
-   |
-   = help: Replace with `trio.lowlevel.checkpoint()`
-
-ℹ Safe fix
-21 21 |     trio.sleep(bar)
-22 22 | 
-23 23 | 
-24    |-trio.sleep(0)  # TRIO115
-   24 |+trio.lowlevel.checkpoint  # TRIO115
-25 25 | 
-26 26 | 
-27 27 | def func():
-
-TRIO115.py:28:14: TRIO115 [*] Use `trio.lowlevel.checkpoint()` instead of `trio.sleep(0)`
-   |
-27 | def func():
-28 |     trio.run(trio.sleep(0))  # TRIO115
-   |              ^^^^^^^^^^^^^ TRIO115
+29 | def func():
+30 |     sleep(0)  # TRIO115
+   |     ^^^^^^^^ TRIO115
    |
    = help: Replace with `trio.lowlevel.checkpoint()`
 
 ℹ Safe fix
+24 24 |     trio.run(trio.sleep(0))  # TRIO115
 25 25 | 
 26 26 | 
-27 27 | def func():
-28    |-    trio.run(trio.sleep(0))  # TRIO115
-   28 |+    trio.run(trio.lowlevel.checkpoint)  # TRIO115
+27    |-from trio import Event, sleep
+   27 |+from trio import Event, sleep, lowlevel
+28 28 | 
+29 29 | def func():
+30    |-    sleep(0)  # TRIO115
+   30 |+    lowlevel.checkpoint  # TRIO115