Detect asyncio.get_running_loop calls in RUF006

astral-sh · Sep 21, 2023 · 35bbae6 · 35bbae6
1 parent b685ee4
commit 35bbae6
Show file tree

Hide file tree

Showing 6 changed files with 103 additions and 60 deletions.
diff --git a/crates/ruff_linter/resources/test/fixtures/ruff/RUF006.py b/crates/ruff_linter/resources/test/fixtures/ruff/RUF006.py
@@ -63,11 +63,28 @@ def f():
     tasks = [asyncio.create_task(task) for task in tasks]
 
 
-# Ok (false negative)
+# OK (false negative)
 def f():
     task = asyncio.create_task(coordinator.ws_connect())
 
 
-# Ok (potential false negative)
+# OK (potential false negative)
 def f():
     do_nothing_with_the_task(asyncio.create_task(coordinator.ws_connect()))
+
+
+# Error
+def f():
+    loop = asyncio.get_running_loop()
+    loop.create_task(coordinator.ws_connect())  # Error
+
+
+# OK
+def f():
+    loop.create_task(coordinator.ws_connect())
+
+
+# OK
+def f():
+    loop = asyncio.get_running_loop()
+    loop.do_thing(coordinator.ws_connect())
diff --git a/crates/ruff_linter/src/rules/flake8_bandit/rules/flask_debug_true.rs b/crates/ruff_linter/src/rules/flake8_bandit/rules/flask_debug_true.rs
@@ -1,7 +1,8 @@
 use ruff_diagnostics::{Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, violation};
 use ruff_python_ast::helpers::is_const_true;
-use ruff_python_ast::{Expr, ExprAttribute, ExprCall, Stmt, StmtAssign};
+use ruff_python_ast::{Expr, ExprAttribute, ExprCall};
+use ruff_python_semantic::analyze::typing;
 use ruff_text_size::Ranged;
 
 use crate::checkers::ast::Checker;
@@ -63,27 +64,11 @@ pub(crate) fn flask_debug_true(checker: &mut Checker, call: &ExprCall) {
         return;
     }
 
-    let Expr::Name(name) = value.as_ref() else {
-        return;
-    };
-
-    if let Some(binding_id) = checker.semantic().resolve_name(name) {
-        if let Some(Stmt::Assign(StmtAssign { value, .. })) = checker
-            .semantic()
-            .binding(binding_id)
-            .statement(checker.semantic())
-        {
-            if let Expr::Call(ExprCall { func, .. }) = value.as_ref() {
-                if checker
-                    .semantic()
-                    .resolve_call_path(func)
-                    .is_some_and(|call_path| matches!(call_path.as_slice(), ["flask", "Flask"]))
-                {
-                    checker
-                        .diagnostics
-                        .push(Diagnostic::new(FlaskDebugTrue, debug_argument.range()));
-                }
-            }
-        }
-    };
+    if typing::resolve_assignment(value, checker.semantic())
+        .is_some_and(|call_path| matches!(call_path.as_slice(), ["flask", "Flask"]))
+    {
+        checker
+            .diagnostics
+            .push(Diagnostic::new(FlaskDebugTrue, debug_argument.range()));
+    }
 }
diff --git a/crates/ruff_linter/src/rules/flake8_bandit/rules/ssh_no_host_key_verification.rs b/crates/ruff_linter/src/rules/flake8_bandit/rules/ssh_no_host_key_verification.rs
@@ -1,6 +1,7 @@
 use ruff_diagnostics::{Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, violation};
-use ruff_python_ast::{Expr, ExprAttribute, ExprCall, Stmt, StmtAssign};
+use ruff_python_ast::{Expr, ExprAttribute, ExprCall};
+use ruff_python_semantic::analyze::typing;
 use ruff_text_size::Ranged;
 
 use crate::checkers::ast::Checker;
@@ -68,30 +69,12 @@ pub(crate) fn ssh_no_host_key_verification(checker: &mut Checker, call: &ExprCal
         return;
     }
 
-    let Expr::Name(name) = value.as_ref() else {
-        return;
-    };
-
-    if let Some(binding_id) = checker.semantic().resolve_name(name) {
-        if let Some(Stmt::Assign(StmtAssign { value, .. })) = checker
-            .semantic()
-            .binding(binding_id)
-            .statement(checker.semantic())
-        {
-            if let Expr::Call(ExprCall { func, .. }) = value.as_ref() {
-                if checker
-                    .semantic()
-                    .resolve_call_path(func)
-                    .is_some_and(|call_path| {
-                        matches!(call_path.as_slice(), ["paramiko", "client", "SSHClient"])
-                    })
-                {
-                    checker.diagnostics.push(Diagnostic::new(
-                        SSHNoHostKeyVerification,
-                        policy_argument.range(),
-                    ));
-                }
-            }
-        }
-    };
+    if typing::resolve_assignment(value, checker.semantic()).is_some_and(|call_path| {
+        matches!(call_path.as_slice(), ["paramiko", "client", "SSHClient"])
+    }) {
+        checker.diagnostics.push(Diagnostic::new(
+            SSHNoHostKeyVerification,
+            policy_argument.range(),
+        ));
+    }
 }
diff --git a/crates/ruff_linter/src/rules/ruff/rules/asyncio_dangling_task.rs b/crates/ruff_linter/src/rules/ruff/rules/asyncio_dangling_task.rs
@@ -4,6 +4,7 @@ use ruff_python_ast::{self as ast, Expr};
 
 use ruff_diagnostics::{Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, violation};
+use ruff_python_semantic::analyze::typing;
 use ruff_text_size::Ranged;
 
 use crate::checkers::ast::Checker;
@@ -70,22 +71,38 @@ pub(crate) fn asyncio_dangling_task(checker: &mut Checker, expr: &Expr) {
         return;
     };
 
-    let Some(method) = checker
+    // Ex) `asyncio.create_task(...)`
+    if let Some(method) = checker
         .semantic()
         .resolve_call_path(func)
         .and_then(|call_path| match call_path.as_slice() {
             ["asyncio", "create_task"] => Some(Method::CreateTask),
             ["asyncio", "ensure_future"] => Some(Method::EnsureFuture),
             _ => None,
         })
-    else {
+    {
+        checker.diagnostics.push(Diagnostic::new(
+            AsyncioDanglingTask { method },
+            expr.range(),
+        ));
         return;
-    };
+    }
 
-    checker.diagnostics.push(Diagnostic::new(
-        AsyncioDanglingTask { method },
-        expr.range(),
-    ));
+    // Ex) `loop = asyncio.get_running_loop(); loop.create_task(...)`
+    if let Expr::Attribute(ast::ExprAttribute { attr, value, .. }) = func.as_ref() {
+        if attr == "create_task" {
+            if typing::resolve_assignment(value, checker.semantic()).is_some_and(|call_path| {
+                matches!(call_path.as_slice(), ["asyncio", "get_running_loop"])
+            }) {
+                checker.diagnostics.push(Diagnostic::new(
+                    AsyncioDanglingTask {
+                        method: Method::CreateTask,
+                    },
+                    expr.range(),
+                ));
+            }
+        }
+    }
 }
 
 #[derive(Debug, PartialEq, Eq, Copy, Clone)]

diff --git a/...ff_linter/src/rules/ruff/snapshots/ruff_linter__rules__ruff__tests__RUF006_RUF006.py.snap b/...ff_linter/src/rules/ruff/snapshots/ruff_linter__rules__ruff__tests__RUF006_RUF006.py.snap
@@ -17,4 +17,12 @@ RUF006.py:11:5: RUF006 Store a reference to the return value of `asyncio.ensure_
    |     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ RUF006
    |
 
+RUF006.py:79:5: RUF006 Store a reference to the return value of `asyncio.create_task`
+   |
+77 | def f():
+78 |     loop = asyncio.get_running_loop()
+79 |     loop.create_task(coordinator.ws_connect())  # Error
+   |     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ RUF006
+   |
+
 
diff --git a/crates/ruff_python_semantic/src/analyze/typing.rs b/crates/ruff_python_semantic/src/analyze/typing.rs
@@ -521,3 +521,36 @@ fn find_parameter<'a>(
         .chain(parameters.kwonlyargs.iter())
         .find(|arg| arg.parameter.name.range() == binding.range())
 }
+
+/// Return the [`CallPath`] of the value to which the given [`Expr`] is assigned, if any.
+///
+/// For example, given:
+/// ```python
+/// import asyncio
+///
+/// loop = asyncio.get_running_loop()
+/// loop.create_task(...)
+/// ```
+///
+/// This function will return `["asyncio", "get_running_loop"]` for the `loop` binding.
+pub fn resolve_assignment<'a>(
+    expr: &'a Expr,
+    semantic: &'a SemanticModel<'a>,
+) -> Option<CallPath<'a>> {
+    let name = expr.as_name_expr()?;
+    let binding_id = semantic.resolve_name(name)?;
+    let statement = semantic.binding(binding_id).statement(semantic)?;
+    match statement {
+        Stmt::Assign(ast::StmtAssign { value, .. }) => {
+            let ast::ExprCall { func, .. } = value.as_call_expr()?;
+            semantic.resolve_call_path(func)
+        }
+        Stmt::AnnAssign(ast::StmtAnnAssign {
+            value: Some(value), ..
+        }) => {
+            let ast::ExprCall { func, .. } = value.as_call_expr()?;
+            semantic.resolve_call_path(func)
+        }
+        _ => None,
+    }
+}