Added support for spawn_local within tokio LocalSets

Author: Andrew J Westlake

pytests/tokio_asyncio/mod.rs

@@ -1,4 +1,4 @@
-use std::time::Duration;
+use std::{rc::Rc, time::Duration};
 
 use pyo3::{prelude::*, wrap_pyfunction};
 
@@ -82,3 +82,22 @@ fn test_init_tokio_twice() -> PyResult<()> {
 
     Ok(())
 }
+
+#[pyo3_asyncio::tokio::test]
+fn test_local_set_coroutine() -> PyResult<()> {
+    tokio::task::LocalSet::new().block_on(pyo3_asyncio::tokio::get_runtime(), async {
+        Python::with_gil(|py| {
+            let non_send_secs = Rc::new(1);
+
+            let py_future = pyo3_asyncio::tokio::into_local_py_future(py, async move {
+                tokio::time::sleep(Duration::from_secs(*non_send_secs)).await;
+                Ok(Python::with_gil(|py| py.None()))
+            })?;
+
+            pyo3_asyncio::into_future(py_future.as_ref(py))
+        })?
+        .await?;
+
+        Ok(())
+    })
+}

src/generic.rs

@@ -23,6 +23,14 @@ pub trait Runtime {
         F: Future<Output = ()> + Send + 'static;
 }
 
+/// Extension trait for async/await runtimes that support spawning local tasks
+pub trait SpawnLocalExt: Runtime {
+    /// Spawn a !Send future onto this runtime's event loop
+    fn spawn_local<F>(fut: F) -> Self::JoinHandle
+    where
+        F: Future<Output = ()> + 'static;
+}
+
 /// Run the event loop until the given Future completes
 ///
 /// After this function returns, the event loop can be resumed with either [`run_until_complete`] or
@@ -236,3 +244,123 @@ where
 
     Ok(future_rx)
 }
+
+/// Convert a `!Send` Rust Future into a Python coroutine with a generic runtime
+///
+/// # Arguments
+/// * `py` - The current PyO3 GIL guard
+/// * `fut` - The Rust future to be converted
+///
+/// # Examples
+///
+/// ```no_run
+/// # use std::{task::{Context, Poll}, pin::Pin, future::Future};
+/// #
+/// # use pyo3_asyncio::generic::{JoinError, SpawnLocalExt, Runtime};
+/// #
+/// # struct MyCustomJoinError;
+/// #
+/// # impl JoinError for MyCustomJoinError {
+/// #     fn is_panic(&self) -> bool {
+/// #         unreachable!()
+/// #     }
+/// # }
+/// #
+/// # struct MyCustomJoinHandle;
+/// #
+/// # impl Future for MyCustomJoinHandle {
+/// #     type Output = Result<(), MyCustomJoinError>;
+/// #
+/// #     fn poll(self: Pin<&mut Self>, _cx: &mut Context) -> Poll<Self::Output> {
+/// #         unreachable!()
+/// #     }
+/// # }
+/// #
+/// # struct MyCustomRuntime;
+/// #
+/// # impl MyCustomRuntime {
+/// #     async fn sleep(_: Duration) {
+/// #         unreachable!()
+/// #     }
+/// # }
+/// #
+/// # impl Runtime for MyCustomRuntime {
+/// #     type JoinError = MyCustomJoinError;
+/// #     type JoinHandle = MyCustomJoinHandle;
+/// #
+/// #     fn spawn<F>(fut: F) -> Self::JoinHandle
+/// #     where
+/// #         F: Future<Output = ()> + Send + 'static
+/// #     {
+/// #         unreachable!()
+/// #     }
+/// # }
+/// #
+/// # impl SpawnLocalExt for MyCustomRuntime {
+/// #     fn spawn_local<F>(fut: F) -> Self::JoinHandle
+/// #     where
+/// #         F: Future<Output = ()> + 'static
+/// #     {
+/// #         unreachable!()
+/// #     }
+/// # }
+/// #
+/// use std::time::Duration;
+///
+/// use pyo3::prelude::*;
+///
+/// /// Awaitable sleep function
+/// #[pyfunction]
+/// fn sleep_for(py: Python, secs: &PyAny) -> PyResult<PyObject> {
+///     let secs = secs.extract()?;
+///
+///     pyo3_asyncio::generic::into_local_py_future::<MyCustomRuntime, _>(py, async move {
+///         MyCustomRuntime::sleep(Duration::from_secs(secs)).await;
+///         Python::with_gil(|py| Ok(py.None()))
+///    })
+/// }
+/// ```
+pub fn into_local_py_future<R, F>(py: Python, fut: F) -> PyResult<PyObject>
+where
+    R: SpawnLocalExt,
+    F: Future<Output = PyResult<PyObject>> + 'static,
+{
+    let future_rx = CREATE_FUTURE.get().expect(EXPECT_INIT).call0(py)?;
+    let future_tx1 = future_rx.clone();
+    let future_tx2 = future_rx.clone();
+
+    R::spawn_local(async move {
+        if let Err(e) = R::spawn_local(async move {
+            let result = fut.await;
+
+            Python::with_gil(move |py| {
+                if set_result(py, future_tx1.as_ref(py), result)
+                    .map_err(dump_err(py))
+                    .is_err()
+                {
+
+                    // Cancelled
+                }
+            });
+        })
+        .await
+        {
+            if e.is_panic() {
+                Python::with_gil(move |py| {
+                    if set_result(
+                        py,
+                        future_tx2.as_ref(py),
+                        Err(PyException::new_err("rust future panicked")),
+                    )
+                    .map_err(dump_err(py))
+                    .is_err()
+                    {
+                        // Cancelled
+                    }
+                });
+            }
+        }
+    });
+
+    Ok(future_rx)
+}

src/tokio.rs

@@ -56,6 +56,15 @@ impl generic::Runtime for TokioRuntime {
     }
 }
 
+impl generic::SpawnLocalExt for TokioRuntime {
+    fn spawn_local<F>(fut: F) -> Self::JoinHandle
+    where
+        F: Future<Output = ()> + 'static,
+    {
+        tokio::task::spawn_local(fut)
+    }
+}
+
 /// Initialize the Tokio Runtime with a custom build
 pub fn init(runtime: Runtime) {
     TOKIO_RUNTIME
@@ -207,3 +216,57 @@ where
 {
     generic::into_coroutine::<TokioRuntime, _>(py, fut)
 }
+
+/// Convert a `!Send` Rust Future into a Python coroutine
+///
+/// # 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<PyObject> {
+///     // Rc is non-send so it cannot be passed into pyo3_asyncio::tokio::into_coroutine
+///     let secs = Rc::new(secs);
+///
+///     pyo3_asyncio::tokio::into_local_py_future(py, async move {
+///         tokio::time::sleep(Duration::from_secs(*secs)).await;
+///         Python::with_gil(|py| Ok(py.None()))
+///     })
+/// }
+///
+/// # #[cfg(all(feature = "tokio-runtime", feature = "attributes"))]
+/// #[pyo3_asyncio::tokio::main]
+/// async fn main() -> PyResult<()> {
+///     // the main coroutine is running in a Send context, so we cannot use LocalSet here. Instead
+///     // we use spawn_blocking in order to use LocalSet::block_on
+///     tokio::task::spawn_blocking(|| {
+///         // LocalSet allows us to work with !Send futures within tokio. Without it, any calls to
+///         // pyo3_asyncio::tokio::into_local_py_future will panic.
+///         tokio::task::LocalSet::new().block_on(pyo3_asyncio::tokio::get_runtime(), async {
+///             Python::with_gil(|py| {
+///                let py_future = sleep_for(py, 1)?;
+///                pyo3_asyncio::into_future(py_future.as_ref(py))
+///             })?
+///             .await?;
+///
+///             Ok(())
+///         })
+///     }).await.unwrap()
+/// }
+/// # #[cfg(not(all(feature = "tokio-runtime", feature = "attributes")))]
+/// # fn main() {}
+/// ```
+pub fn into_local_py_future<F>(py: Python, fut: F) -> PyResult<PyObject>
+where
+    F: Future<Output = PyResult<PyObject>> + 'static,
+{
+    generic::into_local_py_future::<TokioRuntime, _>(py, fut)
+}