Add support for ensure_future for RUF006 (#2943)

README.md

@@ -1502,7 +1502,7 @@ For more, see [tryceratops](https://pypi.org/project/tryceratops/1.1.0/) on PyPI
 | RUF003 | ambiguous-unicode-character-comment | Comment contains ambiguous unicode character `{confusable}` (did you mean `{representant}`?) | 🛠 |
 | RUF004 | keyword-argument-before-star-argument | Keyword argument `{name}` must come after starred arguments |  |
 | RUF005 | unpack-instead-of-concatenating-to-collection-literal | Consider `{expr}` instead of concatenation | 🛠 |
-| RUF006 | [asyncio-dangling-task](https://beta.ruff.rs/docs/rules/asyncio-dangling-task/) | Store a reference to the return value of `asyncio.create_task` |  |
+| RUF006 | [asyncio-dangling-task](https://beta.ruff.rs/docs/rules/asyncio-dangling-task/) | Store a reference to the return value of `asyncio.{method}` |  |
 | RUF100 | unused-noqa | Unused `noqa` directive | 🛠 |
 
 <!-- End auto-generated sections. -->

crates/ruff/resources/test/fixtures/ruff/RUF006.py

@@ -6,6 +6,11 @@ def f():
     asyncio.create_task(coordinator.ws_connect())  # Error
 
 
+# Error
+def f():
+    asyncio.ensure_future(coordinator.ws_connect())  # Error
+
+
 # OK
 def f():
     background_tasks = set()
@@ -22,6 +27,22 @@ def f():
         task.add_done_callback(background_tasks.discard)
 
 
+# OK
+def f():
+    background_tasks = set()
+
+    for i in range(10):
+        task = asyncio.ensure_future(some_coro(param=i))
+
+        # Add task to the set. This creates a strong reference.
+        background_tasks.add(task)
+
+        # To prevent keeping references to finished tasks forever,
+        # make each task remove its own reference from the set after
+        # completion:
+        task.add_done_callback(background_tasks.discard)
+
+
 # OK
 def f():
     ctx.task = asyncio.create_task(make_request())

crates/ruff/src/rules/ruff/rules/asyncio_dangling_task.rs

@@ -1,22 +1,24 @@
 use rustpython_parser::ast::{Expr, ExprKind};
 
-use ruff_macros::{define_violation, derive_message_formats};
-
 use crate::ast::types::{CallPath, Range};
 use crate::registry::Diagnostic;
 use crate::violation::Violation;
+use ruff_macros::{define_violation, derive_message_formats};
+use serde::{Deserialize, Serialize};
+use std::fmt;
 
 define_violation!(
     /// ## What it does
-    /// Checks for `asyncio.create_task` calls that do not store a reference
-    /// to the returned result.
+    /// Checks for `asyncio.create_task` and `asyncio.ensure_future` calls
+    /// that do not store a reference to the returned result.
     ///
     /// ## Why is this bad?
     /// Per the `asyncio` documentation, the event loop only retains a weak
-    /// reference to tasks. If the task returned by `asyncio.create_task` is
-    /// not stored in a variable, or a collection, or otherwise referenced, it
-    /// may be garbage collected at any time. This can lead to unexpected and
-    /// inconsistent behavior, as your tasks may or may not run to completion.
+    /// reference to tasks. If the task returned by `asyncio.create_task` and
+    /// `asyncio.ensure_future` is not stored in a variable, or a collection,
+    /// or otherwise referenced, it may be garbage collected at any time. This
+    /// can lead to unexpected and inconsistent behavior, as your tasks may or
+    /// may not run to completion.
     ///
     /// ## Example
     /// ```python
@@ -49,13 +51,31 @@ define_violation!(
     /// ## References
     /// * [_The Heisenbug lurking in your async code_](https://textual.textualize.io/blog/2023/02/11/the-heisenbug-lurking-in-your-async-code/)
     /// * [`asyncio.create_task`](https://docs.python.org/3/library/asyncio-task.html#asyncio.create_task)
-    pub struct AsyncioDanglingTask;
+    pub struct AsyncioDanglingTask {
+        pub method: Method,
+    }
 );
 
 impl Violation for AsyncioDanglingTask {
     #[derive_message_formats]
     fn message(&self) -> String {
-        format!("Store a reference to the return value of `asyncio.create_task`")
+        let AsyncioDanglingTask { method } = self;
+        format!("Store a reference to the return value of `asyncio.{method}`")
+    }
+}
+
+#[derive(Debug, PartialEq, Eq, Serialize, Deserialize)]
+pub enum Method {
+    CreateTask,
+    EnsureFuture,
+}
+
+impl fmt::Display for Method {
+    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            Method::CreateTask => fmt.write_str("create_task"),
+            Method::EnsureFuture => fmt.write_str("ensure_future"),
+        }
     }
 }
 
@@ -65,14 +85,22 @@ where
     F: FnOnce(&'a Expr) -> Option<CallPath<'a>>,
 {
     if let ExprKind::Call { func, .. } = &expr.node {
-        if resolve_call_path(func).map_or(false, |call_path| {
-            call_path.as_slice() == ["asyncio", "create_task"]
-        }) {
-            return Some(Diagnostic::new(
-                AsyncioDanglingTask,
+        match resolve_call_path(func).as_deref() {
+            Some(["asyncio", "create_task"]) => Some(Diagnostic::new(
+                AsyncioDanglingTask {
+                    method: Method::CreateTask,
+                },
                 Range::from_located(expr),
-            ));
+            )),
+            Some(["asyncio", "ensure_future"]) => Some(Diagnostic::new(
+                AsyncioDanglingTask {
+                    method: Method::EnsureFuture,
+                },
+                Range::from_located(expr),
+            )),
+            _ => None,
         }
+    } else {
+        None
     }
-    None
 }

crates/ruff/src/rules/ruff/snapshots/ruff__rules__ruff__tests__RUF006_RUF006.py.snap

@@ -3,7 +3,8 @@ source: crates/ruff/src/rules/ruff/mod.rs
 expression: diagnostics
 ---
 - kind:
-    AsyncioDanglingTask: ~
+    AsyncioDanglingTask:
+      method: CreateTask
   location:
     row: 6
     column: 4
@@ -12,4 +13,15 @@ expression: diagnostics
     column: 49
   fix: ~
   parent: ~
+- kind:
+    AsyncioDanglingTask:
+      method: EnsureFuture
+  location:
+    row: 11
+    column: 4
+  end_location:
+    row: 11
+    column: 51
+  fix: ~
+  parent: ~