From 2105b9cc1fa68617250c227bfd9fa1ebe5083957 Mon Sep 17 00:00:00 2001
From: Charlie Marsh <charlie.r.marsh@gmail.com>
Date: Tue, 7 Nov 2023 17:19:26 -0500
Subject: [PATCH] Only flag TRIO109 when trio is present

---
 .../resources/test/fixtures/flake8_trio/TRIO109.py   |  3 +++
 .../flake8_trio/rules/async_function_with_timeout.rs | 11 +++++++++++
 ...ules__flake8_trio__tests__TRIO109_TRIO109.py.snap | 12 ++++++------
 crates/ruff_python_semantic/src/model.rs             | 10 ++++++++++
 4 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/crates/ruff_linter/resources/test/fixtures/flake8_trio/TRIO109.py b/crates/ruff_linter/resources/test/fixtures/flake8_trio/TRIO109.py
index 670e0486b2a14..1d7cca20a471b 100644
--- a/crates/ruff_linter/resources/test/fixtures/flake8_trio/TRIO109.py
+++ b/crates/ruff_linter/resources/test/fixtures/flake8_trio/TRIO109.py
@@ -1,3 +1,6 @@
+import trio
+
+
 async def func():
     ...
 
diff --git a/crates/ruff_linter/src/rules/flake8_trio/rules/async_function_with_timeout.rs b/crates/ruff_linter/src/rules/flake8_trio/rules/async_function_with_timeout.rs
index f32421fbd8c39..048be94b03d5a 100644
--- a/crates/ruff_linter/src/rules/flake8_trio/rules/async_function_with_timeout.rs
+++ b/crates/ruff_linter/src/rules/flake8_trio/rules/async_function_with_timeout.rs
@@ -13,6 +13,10 @@ use crate::checkers::ast::Checker;
 /// 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.
+///
 /// ## Example
 /// ```python
 /// async def func():
@@ -40,12 +44,19 @@ pub(crate) fn async_function_with_timeout(
     checker: &mut Checker,
     function_def: &ast::StmtFunctionDef,
 ) {
+    // Detect `async` calls with a `timeout` argument.
     if !function_def.is_async {
         return;
     }
     let Some(timeout) = function_def.parameters.find("timeout") else {
         return;
     };
+
+    // If `trio` isn't in scope, avoid raising the diagnostic.
+    if !checker.semantic().seen(&["trio"]) {
+        return;
+    }
+
     checker.diagnostics.push(Diagnostic::new(
         TrioAsyncFunctionWithTimeout,
         timeout.range(),
diff --git a/crates/ruff_linter/src/rules/flake8_trio/snapshots/ruff_linter__rules__flake8_trio__tests__TRIO109_TRIO109.py.snap b/crates/ruff_linter/src/rules/flake8_trio/snapshots/ruff_linter__rules__flake8_trio__tests__TRIO109_TRIO109.py.snap
index 07df19550bee3..fc08800b672b2 100644
--- a/crates/ruff_linter/src/rules/flake8_trio/snapshots/ruff_linter__rules__flake8_trio__tests__TRIO109_TRIO109.py.snap
+++ b/crates/ruff_linter/src/rules/flake8_trio/snapshots/ruff_linter__rules__flake8_trio__tests__TRIO109_TRIO109.py.snap
@@ -1,18 +1,18 @@
 ---
 source: crates/ruff_linter/src/rules/flake8_trio/mod.rs
 ---
-TRIO109.py:5:16: TRIO109 Prefer `trio.fail_after` and `trio.move_on_after` over manual `async` timeout behavior
+TRIO109.py:8:16: TRIO109 Prefer `trio.fail_after` and `trio.move_on_after` over manual `async` timeout behavior
   |
-5 | async def func(timeout):
+8 | async def func(timeout):
   |                ^^^^^^^ TRIO109
-6 |     ...
+9 |     ...
   |
 
-TRIO109.py:9:16: TRIO109 Prefer `trio.fail_after` and `trio.move_on_after` over manual `async` timeout behavior
+TRIO109.py:12:16: TRIO109 Prefer `trio.fail_after` and `trio.move_on_after` over manual `async` timeout behavior
    |
- 9 | async def func(timeout=10):
+12 | async def func(timeout=10):
    |                ^^^^^^^^^^ TRIO109
-10 |     ...
+13 |     ...
    |
 
 
diff --git a/crates/ruff_python_semantic/src/model.rs b/crates/ruff_python_semantic/src/model.rs
index bd94df3ec47d1..48006dba9e2c3 100644
--- a/crates/ruff_python_semantic/src/model.rs
+++ b/crates/ruff_python_semantic/src/model.rs
@@ -1207,6 +1207,16 @@ impl<'a> SemanticModel<'a> {
         exceptions
     }
 
+    /// Return `true` if the module at the given path was seen anywhere in the semantic model.
+    /// This includes both direct imports (`import trio`) and member imports (`from trio import
+    /// TrioTask`).
+    pub fn seen(&self, module: &[&str]) -> bool {
+        self.bindings
+            .iter()
+            .filter_map(Binding::as_any_import)
+            .any(|import| import.call_path().starts_with(module))
+    }
+
     /// Generate a [`Snapshot`] of the current semantic model.
     pub fn snapshot(&self) -> Snapshot {
         Snapshot {