Merge pull request #11 from awestlake87/bugfix-support-all-awaitables

Bugfix - Missing support for asyncio.Future and asyncio.Task in into_future fn

Author: awestlake87

pytests/common/mod.rs

@@ -29,3 +29,25 @@ pub(super) fn test_blocking_sleep() -> PyResult<()> {
     thread::sleep(Duration::from_secs(1));
     Ok(())
 }
+
+pub(super) async fn test_other_awaitables() -> PyResult<()> {
+    let fut = Python::with_gil(|py| {
+        let functools = py.import("functools")?;
+        let time = py.import("time")?;
+
+        // spawn a blocking sleep in the threadpool executor - returns a task, not a coroutine
+        let task = pyo3_asyncio::get_event_loop(py).call_method1(
+            "run_in_executor",
+            (
+                py.None(),
+                functools.call_method1("partial", (time.getattr("sleep")?, 1))?,
+            ),
+        )?;
+
+        pyo3_asyncio::into_future(task)
+    })?;
+
+    fut.await?;
+
+    Ok(())
+}

pytests/test_async_std_asyncio.rs

@@ -64,6 +64,11 @@ async fn test_into_future() -> PyResult<()> {
     common::test_into_future().await
 }
 
+#[pyo3_asyncio::async_std::test]
+async fn test_other_awaitables() -> PyResult<()> {
+    common::test_other_awaitables().await
+}
+
 #[pyo3_asyncio::async_std::main]
 async fn main() -> pyo3::PyResult<()> {
     pyo3_asyncio::testing::main().await

pytests/tokio_asyncio/mod.rs

@@ -62,3 +62,8 @@ fn test_blocking_sleep() -> PyResult<()> {
 async fn test_into_future() -> PyResult<()> {
     common::test_into_future().await
 }
+
+#[pyo3_asyncio::tokio::test]
+async fn test_other_awaitables() -> PyResult<()> {
+    common::test_other_awaitables().await
+}

src/lib.rs

@@ -141,12 +141,16 @@ pub mod doc_test {
 const EXPECT_INIT: &str = "PyO3 Asyncio has not been initialized";
 
 static ASYNCIO: OnceCell<PyObject> = OnceCell::new();
+static ENSURE_FUTURE: OnceCell<PyObject> = OnceCell::new();
 static EVENT_LOOP: OnceCell<PyObject> = OnceCell::new();
 static EXECUTOR: OnceCell<PyObject> = OnceCell::new();
 static CALL_SOON: OnceCell<PyObject> = OnceCell::new();
-static CREATE_TASK: OnceCell<PyObject> = OnceCell::new();
 static CREATE_FUTURE: OnceCell<PyObject> = OnceCell::new();
 
+fn ensure_future(py: Python) -> &PyAny {
+    ENSURE_FUTURE.get().expect(EXPECT_INIT).as_ref(py)
+}
+
 #[allow(clippy::needless_doctest_main)]
 /// Wraps the provided function with the initialization and finalization for PyO3 Asyncio
 ///
@@ -192,6 +196,9 @@ where
 /// Must be called at the start of your program
 fn try_init(py: Python) -> PyResult<()> {
     let asyncio = py.import("asyncio")?;
+
+    let ensure_future = asyncio.getattr("ensure_future")?;
+
     let event_loop = asyncio.call_method0("get_event_loop")?;
     let executor = py
         .import("concurrent.futures.thread")?
@@ -201,14 +208,13 @@ fn try_init(py: Python) -> PyResult<()> {
     event_loop.call_method1("set_default_executor", (executor,))?;
 
     let call_soon = event_loop.getattr("call_soon_threadsafe")?;
-    let create_task = asyncio.getattr("run_coroutine_threadsafe")?;
     let create_future = event_loop.getattr("create_future")?;
 
     ASYNCIO.get_or_init(|| asyncio.into());
+    ENSURE_FUTURE.get_or_init(|| ensure_future.into());
     EVENT_LOOP.get_or_init(|| event_loop.into());
     EXECUTOR.get_or_init(|| executor.into());
     CALL_SOON.get_or_init(|| call_soon.into());
-    CREATE_TASK.get_or_init(|| create_task.into());
     CREATE_FUTURE.get_or_init(|| create_future.into());
 
     Ok(())
@@ -321,6 +327,26 @@ impl PyTaskCompleter {
     }
 }
 
+#[pyclass]
+struct PyEnsureFuture {
+    awaitable: PyObject,
+    tx: Option<oneshot::Sender<PyResult<PyObject>>>,
+}
+
+#[pymethods]
+impl PyEnsureFuture {
+    #[call]
+    pub fn __call__(&mut self) -> PyResult<()> {
+        Python::with_gil(|py| {
+            let task = ensure_future(py).call1((self.awaitable.as_ref(py),))?;
+            let on_complete = PyTaskCompleter { tx: self.tx.take() };
+            task.call_method1("add_done_callback", (on_complete,))?;
+
+            Ok(())
+        })
+    }
+}
+
 /// Convert a Python `awaitable` into a Rust Future
 ///
 /// This function converts the `awaitable` into a Python Task using `run_coroutine_threadsafe`. A
@@ -373,13 +399,13 @@ pub fn into_future(awaitable: &PyAny) -> PyResult<impl Future<Output = PyResult<
     let py = awaitable.py();
     let (tx, rx) = oneshot::channel();
 
-    let task = CREATE_TASK
-        .get()
-        .expect(EXPECT_INIT)
-        .call1(py, (awaitable, get_event_loop(py)))?;
-    let on_complete = PyTaskCompleter { tx: Some(tx) };
-
-    task.call_method1(py, "add_done_callback", (on_complete,))?;
+    CALL_SOON.get().expect(EXPECT_INIT).call1(
+        py,
+        (PyEnsureFuture {
+            awaitable: awaitable.into(),
+            tx: Some(tx),
+        },),
+    )?;
 
     Ok(async move {
         match rx.await {

src/tokio.rs

@@ -88,7 +88,7 @@ pub fn init_multi_thread() {
         Builder::new_multi_thread()
             .enable_all()
             .build()
-            .expect("Couldn't build the current-thread Tokio runtime"),
+            .expect("Couldn't build the multi-thread Tokio runtime"),
     );
 }