Added a 'cancellable' set of functions to allow Rust futures to be cancelled from Python (may become default behaviour in next release)

Author: Andrew J Westlake

Cargo.toml

@@ -106,6 +106,7 @@ clap = { version = "2.33", optional = true }
 futures = "0.3"
 inventory = "0.1"
 once_cell = "1.5"
+pin-project-lite = "0.2"
 pyo3 = "0.14"
 pyo3-asyncio-macros = { path = "pyo3-asyncio-macros", version = "=0.14.0", optional = true }
 

pytests/test_async_std_asyncio.rs

@@ -1,6 +1,10 @@
 mod common;
 
-use std::{rc::Rc, time::Duration};
+use std::{
+    rc::Rc,
+    sync::{Arc, Mutex},
+    time::Duration,
+};
 
 use async_std::task;
 use pyo3::{
@@ -171,12 +175,19 @@ async fn test_local_future_into_py() -> PyResult<()> {
 
 #[pyo3_asyncio::async_std::test]
 async fn test_cancel() -> PyResult<()> {
+    let completed = Arc::new(Mutex::new(false));
+
     let py_future = Python::with_gil(|py| -> PyResult<PyObject> {
-        Ok(pyo3_asyncio::async_std::future_into_py(py, async {
-            async_std::task::sleep(Duration::from_secs(1)).await;
-            Ok(Python::with_gil(|py| py.None()))
-        })?
-        .into())
+        let completed = Arc::clone(&completed);
+        Ok(
+            pyo3_asyncio::async_std::cancellable_future_into_py(py, async move {
+                async_std::task::sleep(Duration::from_secs(1)).await;
+                *completed.lock().unwrap() = true;
+
+                Ok(Python::with_gil(|py| py.None()))
+            })?
+            .into(),
+        )
     })?;
 
     if let Err(e) = Python::with_gil(|py| -> PyResult<_> {
@@ -198,6 +209,11 @@ async fn test_cancel() -> PyResult<()> {
         panic!("expected CancelledError");
     }
 
+    async_std::task::sleep(Duration::from_secs(1)).await;
+    if *completed.lock().unwrap() {
+        panic!("future still completed")
+    }
+
     Ok(())
 }
 

pytests/tokio_asyncio/mod.rs

@@ -1,4 +1,8 @@
-use std::{rc::Rc, time::Duration};
+use std::{
+    rc::Rc,
+    sync::{Arc, Mutex},
+    time::Duration,
+};
 
 use pyo3::{
     prelude::*,
@@ -168,12 +172,19 @@ async fn test_panic() -> PyResult<()> {
 
 #[pyo3_asyncio::tokio::test]
 async fn test_cancel() -> PyResult<()> {
+    let completed = Arc::new(Mutex::new(false));
+
     let py_future = Python::with_gil(|py| -> PyResult<PyObject> {
-        Ok(pyo3_asyncio::tokio::future_into_py(py, async {
-            tokio::time::sleep(Duration::from_secs(1)).await;
-            Ok(Python::with_gil(|py| py.None()))
-        })?
-        .into())
+        let completed = Arc::clone(&completed);
+        Ok(
+            pyo3_asyncio::tokio::cancellable_future_into_py(py, async move {
+                tokio::time::sleep(Duration::from_secs(1)).await;
+                *completed.lock().unwrap() = true;
+
+                Ok(Python::with_gil(|py| py.None()))
+            })?
+            .into(),
+        )
     })?;
 
     if let Err(e) = Python::with_gil(|py| -> PyResult<_> {
@@ -195,8 +206,14 @@ async fn test_cancel() -> PyResult<()> {
         panic!("expected CancelledError");
     }
 
+    tokio::time::sleep(Duration::from_secs(1)).await;
+    if *completed.lock().unwrap() {
+        panic!("future still completed")
+    }
+
     Ok(())
 }
+
 /// This module is implemented in Rust.
 #[pymodule]
 fn test_mod(_py: Python, m: &PyModule) -> PyResult<()> {

src/async_std.rs

@@ -273,6 +273,42 @@ where
     generic::future_into_py_with_loop::<AsyncStdRuntime, F>(event_loop, fut)
 }
 
+/// Convert a Rust Future into a Python awaitable
+///
+/// Unlike [`future_into_py_with_loop`], this function will stop the Rust future from running when
+/// the `asyncio.Future` is cancelled from Python.
+///
+/// # Arguments
+/// * `event_loop` - The Python event loop that the awaitable should be attached to
+/// * `fut` - The Rust future to be converted
+///
+/// # Examples
+///
+/// ```
+/// use std::time::Duration;
+///
+/// use pyo3::prelude::*;
+///
+/// /// Awaitable sleep function
+/// #[pyfunction]
+/// fn sleep_for<'p>(py: Python<'p>, secs: &'p PyAny) -> PyResult<&'p PyAny> {
+///     let secs = secs.extract()?;
+///     pyo3_asyncio::async_std::cancellable_future_into_py_with_loop(
+///         pyo3_asyncio::async_std::get_current_loop(py)?,
+///         async move {
+///             async_std::task::sleep(Duration::from_secs(secs)).await;
+///             Python::with_gil(|py| Ok(py.None()))
+///         }
+///     )
+/// }
+/// ```
+pub fn cancellable_future_into_py_with_loop<F>(event_loop: &PyAny, fut: F) -> PyResult<&PyAny>
+where
+    F: Future<Output = PyResult<PyObject>> + Send + 'static,
+{
+    generic::cancellable_future_into_py_with_loop::<AsyncStdRuntime, F>(event_loop, fut)
+}
+
 /// Convert a Rust Future into a Python awaitable
 ///
 /// # Arguments
@@ -303,6 +339,39 @@ where
     generic::future_into_py::<AsyncStdRuntime, _>(py, fut)
 }
 
+/// Convert a Rust Future into a Python awaitable
+///
+/// Unlike [`future_into_py`], this function will stop the Rust future from running when
+/// the `asyncio.Future` is cancelled from Python.
+///
+/// # Arguments
+/// * `py` - The current PyO3 GIL guard
+/// * `fut` - The Rust future to be converted
+///
+/// # Examples
+///
+/// ```
+/// use std::time::Duration;
+///
+/// use pyo3::prelude::*;
+///
+/// /// Awaitable sleep function
+/// #[pyfunction]
+/// fn sleep_for<'p>(py: Python<'p>, secs: &'p PyAny) -> PyResult<&'p PyAny> {
+///     let secs = secs.extract()?;
+///     pyo3_asyncio::async_std::cancellable_future_into_py(py, async move {
+///         async_std::task::sleep(Duration::from_secs(secs)).await;
+///         Python::with_gil(|py| Ok(py.None()))
+///     })
+/// }
+/// ```
+pub fn cancellable_future_into_py<F>(py: Python, fut: F) -> PyResult<&PyAny>
+where
+    F: Future<Output = PyResult<PyObject>> + Send + 'static,
+{
+    generic::cancellable_future_into_py::<AsyncStdRuntime, _>(py, fut)
+}
+
 /// Convert a `!Send` Rust Future into a Python awaitable
 ///
 /// # Arguments