Added missing local_cancellable_* variants to tokio and async-std runtimes and added new tests for them

Author: Andrew J Westlake

pytests/test_async_std_asyncio.rs

@@ -217,6 +217,52 @@ async fn test_cancel() -> PyResult<()> {
     Ok(())
 }
 
+#[pyo3_asyncio::async_std::test]
+fn test_local_cancel(event_loop: PyObject) -> PyResult<()> {
+    async_std::task::block_on(pyo3_asyncio::async_std::scope_local(event_loop, async {
+        let completed = Arc::new(Mutex::new(false));
+
+        let py_future = Python::with_gil(|py| -> PyResult<PyObject> {
+            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<_> {
+            py_future.as_ref(py).call_method0("cancel")?;
+            pyo3_asyncio::async_std::into_future(py_future.as_ref(py))
+        })?
+        .await
+        {
+            Python::with_gil(|py| -> PyResult<()> {
+                assert!(py
+                    .import("asyncio")?
+                    .getattr("CancelledError")?
+                    .downcast::<PyType>()
+                    .unwrap()
+                    .is_instance(e.pvalue(py))?);
+                Ok(())
+            })?;
+        } else {
+            panic!("expected CancelledError");
+        }
+
+        async_std::task::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<()> {

pytests/tokio_asyncio/mod.rs

@@ -214,6 +214,54 @@ async fn test_cancel() -> PyResult<()> {
     Ok(())
 }
 
+#[pyo3_asyncio::tokio::test]
+fn test_local_cancel(event_loop: PyObject) -> PyResult<()> {
+    tokio::task::LocalSet::new().block_on(
+        pyo3_asyncio::tokio::get_runtime(),
+        pyo3_asyncio::tokio::scope_local(event_loop, async {
+            let completed = Arc::new(Mutex::new(false));
+            let py_future = Python::with_gil(|py| -> PyResult<PyObject> {
+                let completed = Arc::clone(&completed);
+                Ok(
+                    pyo3_asyncio::tokio::local_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<_> {
+                py_future.as_ref(py).call_method0("cancel")?;
+                pyo3_asyncio::tokio::into_future(py_future.as_ref(py))
+            })?
+            .await
+            {
+                Python::with_gil(|py| -> PyResult<()> {
+                    assert!(py
+                        .import("asyncio")?
+                        .getattr("CancelledError")?
+                        .downcast::<PyType>()
+                        .unwrap()
+                        .is_instance(e.pvalue(py))?);
+                    Ok(())
+                })?;
+            } else {
+                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

@@ -278,8 +278,8 @@ where
 /// Unlike [`future_into_py_with_loop`], this function will stop the Rust future from running when
 /// the `asyncio.Future` is cancelled from Python.
 ///
-/// __This function will be deprecated in favor of [`future_into_py_with_loop`] in `v0.15` because 
-/// it will become the default behaviour. In `v0.15`, any calls to this function can be seamlessly 
+/// __This function will be deprecated in favor of [`future_into_py_with_loop`] in `v0.15` because
+/// it will become the default behaviour. In `v0.15`, any calls to this function can be seamlessly
 /// replaced with [`future_into_py_with_loop`].__
 ///
 /// # Arguments
@@ -348,8 +348,8 @@ where
 /// Unlike [`future_into_py`], this function will stop the Rust future from running when
 /// the `asyncio.Future` is cancelled from Python.
 ///
-/// __This function will be deprecated in favor of [`future_into_py`] in `v0.15` because 
-/// it will become the default behaviour. In `v0.15`, any calls to this function can be seamlessly 
+/// __This function will be deprecated in favor of [`future_into_py`] in `v0.15` because
+/// it will become the default behaviour. In `v0.15`, any calls to this function can be seamlessly
 /// replaced with [`future_into_py`].__
 ///
 /// # Arguments
@@ -428,6 +428,61 @@ where
     generic::local_future_into_py_with_loop::<AsyncStdRuntime, _>(event_loop, fut)
 }
 
+/// Convert a `!Send` Rust Future into a Python awaitable
+///
+/// Unlike [`local_future_into_py_with_loop`], this function will stop the Rust future from running when
+/// the `asyncio.Future` is cancelled from Python.
+///
+/// __This function will be deprecated in favor of [`local_future_into_py_with_loop`] in `v0.15` because
+/// it will become the default behaviour. In `v0.15`, any calls to this function can be seamlessly
+/// replaced with [`local_future_into_py_with_loop`].__
+///
+/// # Arguments
+/// * `event_loop` - The Python event loop that the awaitable should be attached to
+/// * `fut` - The Rust future to be converted
+///
+/// # Examples
+///
+/// ```
+/// use std::{rc::Rc, time::Duration};
+///
+/// use pyo3::prelude::*;
+///
+/// /// Awaitable non-send sleep function
+/// #[pyfunction]
+/// fn sleep_for(py: Python, secs: u64) -> PyResult<&PyAny> {
+///     // Rc is non-send so it cannot be passed into pyo3_asyncio::async_std::future_into_py
+///     let secs = Rc::new(secs);
+///     Ok(pyo3_asyncio::async_std::local_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()))
+///         }
+///     )?.into())
+/// }
+///
+/// # #[cfg(all(feature = "async-std-runtime", feature = "attributes"))]
+/// #[pyo3_asyncio::async_std::main]
+/// async fn main() -> PyResult<()> {
+///     Python::with_gil(|py| {
+///         let py_future = sleep_for(py, 1)?;
+///         pyo3_asyncio::async_std::into_future(py_future)
+///     })?
+///     .await?;
+///
+///     Ok(())
+/// }
+/// # #[cfg(not(all(feature = "async-std-runtime", feature = "attributes")))]
+/// # fn main() {}
+/// ```
+pub fn local_cancellable_future_into_py_with_loop<F>(event_loop: &PyAny, fut: F) -> PyResult<&PyAny>
+where
+    F: Future<Output = PyResult<PyObject>> + 'static,
+{
+    generic::local_cancellable_future_into_py_with_loop::<AsyncStdRuntime, _>(event_loop, fut)
+}
+
 /// Convert a `!Send` Rust Future into a Python awaitable
 ///
 /// # Arguments
@@ -473,6 +528,58 @@ where
     generic::local_future_into_py::<AsyncStdRuntime, _>(py, fut)
 }
 
+/// Convert a `!Send` Rust Future into a Python awaitable
+///
+/// Unlike [`local_future_into_py`], this function will stop the Rust future from running when
+/// the `asyncio.Future` is cancelled from Python.
+///
+/// __This function will be deprecated in favor of [`local_future_into_py`] in `v0.15` because
+/// it will become the default behaviour. In `v0.15`, any calls to this function can be seamlessly
+/// replaced with [`local_future_into_py`].__
+///
+/// # Arguments
+/// * `py` - The current PyO3 GIL guard
+/// * `fut` - The Rust future to be converted
+///
+/// # Examples
+///
+/// ```
+/// use std::{rc::Rc, time::Duration};
+///
+/// use pyo3::prelude::*;
+///
+/// /// Awaitable non-send sleep function
+/// #[pyfunction]
+/// fn sleep_for(py: Python, secs: u64) -> PyResult<&PyAny> {
+///     // Rc is non-send so it cannot be passed into pyo3_asyncio::async_std::future_into_py
+///     let secs = Rc::new(secs);
+///     pyo3_asyncio::async_std::local_cancellable_future_into_py(py, async move {
+///         async_std::task::sleep(Duration::from_secs(*secs)).await;
+///         Python::with_gil(|py| Ok(py.None()))
+///     })
+/// }
+///
+/// # #[cfg(all(feature = "async-std-runtime", feature = "attributes"))]
+/// #[pyo3_asyncio::async_std::main]
+/// async fn main() -> PyResult<()> {
+///     Python::with_gil(|py| {
+///         let py_future = sleep_for(py, 1)?;
+///         pyo3_asyncio::async_std::into_future(py_future)
+///     })?
+///     .await?;
+///
+///     Ok(())
+/// }
+/// # #[cfg(not(all(feature = "async-std-runtime", feature = "attributes")))]
+/// # fn main() {}
+/// ```
+pub fn local_cancellable_future_into_py<F>(py: Python, fut: F) -> PyResult<&PyAny>
+where
+    F: Future<Output = PyResult<PyObject>> + 'static,
+{
+    generic::local_cancellable_future_into_py::<AsyncStdRuntime, _>(py, fut)
+}
+
 /// Convert a Python `awaitable` into a Rust Future
 ///
 /// This function converts the `awaitable` into a Python Task using `run_coroutine_threadsafe`. A