pyo3_async_runtimes/
generic.rs

1//! Generic implementations of PyO3 Asyncio utilities that can be used for any Rust runtime
2//!
3//! Items marked with
4//! <span
5//!   class="module-item stab portability"
6//!   style="display: inline; border-radius: 3px; padding: 2px; font-size: 80%; line-height: 1.2;"
7//! ><code>unstable-streams</code></span>
8//! > are only available when the `unstable-streams` Cargo feature is enabled:
9//!
10//! ```toml
11//! [dependencies.pyo3-async-runtimes]
12//! version = "0.24"
13//! features = ["unstable-streams"]
14//! ```
15
16use std::{
17    future::Future,
18    pin::Pin,
19    sync::{Arc, Mutex},
20    task::{Context, Poll},
21};
22
23use crate::{
24    asyncio, call_soon_threadsafe, close, create_future, dump_err, err::RustPanic,
25    get_running_loop, into_future_with_locals, TaskLocals,
26};
27#[cfg(feature = "unstable-streams")]
28use futures_channel::mpsc;
29use futures_channel::oneshot;
30#[cfg(feature = "unstable-streams")]
31use futures_util::SinkExt;
32use pin_project_lite::pin_project;
33use pyo3::prelude::*;
34use pyo3::IntoPyObjectExt;
35#[cfg(feature = "unstable-streams")]
36use std::marker::PhantomData;
37
38/// Generic utilities for a JoinError
39pub trait JoinError {
40    /// Check if the spawned task exited because of a panic
41    fn is_panic(&self) -> bool;
42    /// Get the panic object associated with the error.  Panics if `is_panic` is not true.
43    fn into_panic(self) -> Box<dyn std::any::Any + Send + 'static>;
44}
45
46/// Generic Rust async/await runtime
47pub trait Runtime: Send + 'static {
48    /// The error returned by a JoinHandle after being awaited
49    type JoinError: JoinError + Send;
50    /// A future that completes with the result of the spawned task
51    type JoinHandle: Future<Output = Result<(), Self::JoinError>> + Send;
52
53    /// Spawn a future onto this runtime's event loop
54    fn spawn<F>(fut: F) -> Self::JoinHandle
55    where
56        F: Future<Output = ()> + Send + 'static;
57
58    /// Spawn a function onto this runtime's blocking event loop
59    fn spawn_blocking<F>(f: F) -> Self::JoinHandle
60    where
61        F: FnOnce() + Send + 'static;
62}
63
64/// Extension trait for async/await runtimes that support spawning local tasks
65pub trait SpawnLocalExt: Runtime {
66    /// Spawn a !Send future onto this runtime's event loop
67    fn spawn_local<F>(fut: F) -> Self::JoinHandle
68    where
69        F: Future<Output = ()> + 'static;
70}
71
72/// Exposes the utilities necessary for using task-local data in the Runtime
73pub trait ContextExt: Runtime {
74    /// Set the task locals for the given future
75    fn scope<F, R>(locals: TaskLocals, fut: F) -> Pin<Box<dyn Future<Output = R> + Send>>
76    where
77        F: Future<Output = R> + Send + 'static;
78
79    /// Get the task locals for the current task
80    fn get_task_locals() -> Option<TaskLocals>;
81}
82
83/// Adds the ability to scope task-local data for !Send futures
84pub trait LocalContextExt: Runtime {
85    /// Set the task locals for the given !Send future
86    fn scope_local<F, R>(locals: TaskLocals, fut: F) -> Pin<Box<dyn Future<Output = R>>>
87    where
88        F: Future<Output = R> + 'static;
89}
90
91/// Get the current event loop from either Python or Rust async task local context
92///
93/// This function first checks if the runtime has a task-local reference to the Python event loop.
94/// If not, it calls [`get_running_loop`](crate::get_running_loop`) to get the event loop associated
95/// with the current OS thread.
96pub fn get_current_loop<R>(py: Python) -> PyResult<Bound<PyAny>>
97where
98    R: ContextExt,
99{
100    if let Some(locals) = R::get_task_locals() {
101        Ok(locals.0.event_loop.clone_ref(py).into_bound(py))
102    } else {
103        get_running_loop(py)
104    }
105}
106
107/// Either copy the task locals from the current task OR get the current running loop and
108/// contextvars from Python.
109pub fn get_current_locals<R>(py: Python) -> PyResult<TaskLocals>
110where
111    R: ContextExt,
112{
113    if let Some(locals) = R::get_task_locals() {
114        Ok(locals)
115    } else {
116        Ok(TaskLocals::with_running_loop(py)?.copy_context(py)?)
117    }
118}
119
120/// Run the event loop until the given Future completes
121///
122/// After this function returns, the event loop can be resumed with [`run_until_complete`]
123///
124/// # Arguments
125/// * `event_loop` - The Python event loop that should run the future
126/// * `fut` - The future to drive to completion
127///
128/// # Examples
129///
130/// ```no_run
131/// # use std::{any::Any, task::{Context, Poll}, pin::Pin, future::Future};
132/// #
133/// # use pyo3_async_runtimes::{
134/// #     TaskLocals,
135/// #     generic::{JoinError, SpawnLocalExt, ContextExt, LocalContextExt, Runtime}
136/// # };
137/// #
138/// # struct MyCustomJoinError;
139/// #
140/// # impl JoinError for MyCustomJoinError {
141/// #     fn is_panic(&self) -> bool {
142/// #         unreachable!()
143/// #     }
144/// #     fn into_panic(self) -> Box<(dyn Any + Send + 'static)> {
145/// #         unreachable!()
146/// #     }
147/// # }
148/// #
149/// # struct MyCustomJoinHandle;
150/// #
151/// # impl Future for MyCustomJoinHandle {
152/// #     type Output = Result<(), MyCustomJoinError>;
153/// #
154/// #     fn poll(self: Pin<&mut Self>, _cx: &mut Context) -> Poll<Self::Output> {
155/// #         unreachable!()
156/// #     }
157/// # }
158/// #
159/// # struct MyCustomRuntime;
160/// #
161/// # impl Runtime for MyCustomRuntime {
162/// #     type JoinError = MyCustomJoinError;
163/// #     type JoinHandle = MyCustomJoinHandle;
164/// #
165/// #     fn spawn<F>(fut: F) -> Self::JoinHandle
166/// #     where
167/// #         F: Future<Output = ()> + Send + 'static
168/// #     {
169/// #         unreachable!()
170/// #     }
171/// #
172/// #     fn spawn_blocking<F>(f: F) -> Self::JoinHandle where F: FnOnce() + Send + 'static {
173/// #         unreachable!()
174/// #     }
175/// # }
176/// #
177/// # impl ContextExt for MyCustomRuntime {
178/// #     fn scope<F, R>(locals: TaskLocals, fut: F) -> Pin<Box<dyn Future<Output = R> + Send>>
179/// #     where
180/// #         F: Future<Output = R> + Send + 'static
181/// #     {
182/// #         unreachable!()
183/// #     }
184/// #     fn get_task_locals() -> Option<TaskLocals> {
185/// #         unreachable!()
186/// #     }
187/// # }
188/// #
189/// # use std::time::Duration;
190/// #
191/// # use pyo3::prelude::*;
192/// #
193/// # Python::attach(|py| -> PyResult<()> {
194/// # let event_loop = py.import("asyncio")?.call_method0("new_event_loop")?;
195/// # #[cfg(feature = "tokio-runtime")]
196/// pyo3_async_runtimes::generic::run_until_complete::<MyCustomRuntime, _, _>(&event_loop, async move {
197///     tokio::time::sleep(Duration::from_secs(1)).await;
198///     Ok(())
199/// })?;
200/// # Ok(())
201/// # }).unwrap();
202/// ```
203pub fn run_until_complete<R, F, T>(event_loop: &Bound<PyAny>, fut: F) -> PyResult<T>
204where
205    R: Runtime + ContextExt,
206    F: Future<Output = PyResult<T>> + Send + 'static,
207    T: Send + Sync + 'static,
208{
209    let py = event_loop.py();
210    let result_tx = Arc::new(Mutex::new(None));
211    let result_rx = Arc::clone(&result_tx);
212    let coro = future_into_py_with_locals::<R, _, ()>(
213        py,
214        TaskLocals::new(event_loop.clone()).copy_context(py)?,
215        async move {
216            let val = fut.await?;
217            if let Ok(mut result) = result_tx.lock() {
218                *result = Some(val);
219            }
220            Ok(())
221        },
222    )?;
223
224    event_loop.call_method1(pyo3::intern!(py, "run_until_complete"), (coro,))?;
225
226    let result = result_rx.lock().unwrap().take().unwrap();
227    Ok(result)
228}
229
230/// Run the event loop until the given Future completes
231///
232/// # Arguments
233/// * `py` - The current PyO3 GIL guard
234/// * `fut` - The future to drive to completion
235///
236/// # Examples
237///
238/// ```no_run
239/// # use std::{any::Any, task::{Context, Poll}, pin::Pin, future::Future};
240/// #
241/// # use pyo3_async_runtimes::{
242/// #     TaskLocals,
243/// #     generic::{JoinError, SpawnLocalExt, ContextExt, LocalContextExt, Runtime}
244/// # };
245/// #
246/// # struct MyCustomJoinError;
247/// #
248/// # impl JoinError for MyCustomJoinError {
249/// #     fn is_panic(&self) -> bool {
250/// #         unreachable!()
251/// #     }
252/// #     fn into_panic(self) -> Box<(dyn Any + Send + 'static)> {
253/// #         unreachable!()
254/// #     }
255/// # }
256/// #
257/// # struct MyCustomJoinHandle;
258/// #
259/// # impl Future for MyCustomJoinHandle {
260/// #     type Output = Result<(), MyCustomJoinError>;
261/// #
262/// #     fn poll(self: Pin<&mut Self>, _cx: &mut Context) -> Poll<Self::Output> {
263/// #         unreachable!()
264/// #     }
265/// # }
266/// #
267/// # struct MyCustomRuntime;
268/// #
269/// # impl Runtime for MyCustomRuntime {
270/// #     type JoinError = MyCustomJoinError;
271/// #     type JoinHandle = MyCustomJoinHandle;
272/// #
273/// #     fn spawn<F>(fut: F) -> Self::JoinHandle
274/// #     where
275/// #         F: Future<Output = ()> + Send + 'static
276/// #     {
277/// #         unreachable!()
278/// #     }
279/// #
280/// #     fn spawn_blocking<F>(f: F) -> Self::JoinHandle where F: FnOnce() + Send + 'static {
281/// #         unreachable!()
282/// #     }
283/// # }
284/// #
285/// # impl ContextExt for MyCustomRuntime {
286/// #     fn scope<F, R>(locals: TaskLocals, fut: F) -> Pin<Box<dyn Future<Output = R> + Send>>
287/// #     where
288/// #         F: Future<Output = R> + Send + 'static
289/// #     {
290/// #         unreachable!()
291/// #     }
292/// #     fn get_task_locals() -> Option<TaskLocals> {
293/// #         unreachable!()
294/// #     }
295/// # }
296/// #
297/// # use std::time::Duration;
298/// # async fn custom_sleep(_duration: Duration) { }
299/// #
300/// # use pyo3::prelude::*;
301/// #
302/// fn main() {
303///     Python::attach(|py| {
304///         pyo3_async_runtimes::generic::run::<MyCustomRuntime, _, _>(py, async move {
305///             custom_sleep(Duration::from_secs(1)).await;
306///             Ok(())
307///         })
308///         .map_err(|e| {
309///             e.print_and_set_sys_last_vars(py);
310///         })
311///         .unwrap();
312///     })
313/// }
314/// ```
315pub fn run<R, F, T>(py: Python, fut: F) -> PyResult<T>
316where
317    R: Runtime + ContextExt,
318    F: Future<Output = PyResult<T>> + Send + 'static,
319    T: Send + Sync + 'static,
320{
321    let event_loop = asyncio(py)?.call_method0(pyo3::intern!(py, "new_event_loop"))?;
322
323    let result = run_until_complete::<R, F, T>(&event_loop, fut);
324
325    close(event_loop)?;
326
327    result
328}
329
330fn cancelled(future: &Bound<PyAny>) -> PyResult<bool> {
331    future
332        .getattr(pyo3::intern!(future.py(), "cancelled"))?
333        .call0()?
334        .is_truthy()
335}
336
337#[pyclass]
338struct CheckedCompletor;
339
340#[pymethods]
341impl CheckedCompletor {
342    fn __call__(
343        &self,
344        future: &Bound<PyAny>,
345        complete: &Bound<PyAny>,
346        value: &Bound<PyAny>,
347    ) -> PyResult<()> {
348        if cancelled(future)? {
349            return Ok(());
350        }
351
352        complete.call1((value,))?;
353
354        Ok(())
355    }
356}
357
358fn set_result(
359    event_loop: &Bound<PyAny>,
360    future: &Bound<PyAny>,
361    result: PyResult<Py<PyAny>>,
362) -> PyResult<()> {
363    let py = event_loop.py();
364    let none = py.None().into_bound(py);
365
366    let (complete, val) = match result {
367        Ok(val) => (
368            future.getattr(pyo3::intern!(py, "set_result"))?,
369            val.into_pyobject(py)?,
370        ),
371        Err(err) => (
372            future.getattr(pyo3::intern!(py, "set_exception"))?,
373            err.into_bound_py_any(py)?,
374        ),
375    };
376    call_soon_threadsafe(event_loop, &none, (CheckedCompletor, future, complete, val))?;
377
378    Ok(())
379}
380
381/// Convert a Python `awaitable` into a Rust Future
382///
383/// This function simply forwards the future and the task locals returned by [`get_current_locals`]
384/// to [`into_future_with_locals`](`crate::into_future_with_locals`). See
385/// [`into_future_with_locals`](`crate::into_future_with_locals`) for more details.
386///
387/// # Arguments
388/// * `awaitable` - The Python `awaitable` to be converted
389///
390/// # Examples
391///
392/// ```no_run
393/// # use std::{any::Any, pin::Pin, future::Future, task::{Context, Poll}, time::Duration};
394/// # use std::ffi::CString;
395/// #
396/// # use pyo3::prelude::*;
397/// #
398/// # use pyo3_async_runtimes::{
399/// #     TaskLocals,
400/// #     generic::{JoinError, SpawnLocalExt, ContextExt, LocalContextExt, Runtime}
401/// # };
402/// #
403/// # struct MyCustomJoinError;
404/// #
405/// # impl JoinError for MyCustomJoinError {
406/// #     fn is_panic(&self) -> bool {
407/// #         unreachable!()
408/// #     }
409/// #     fn into_panic(self) -> Box<(dyn Any + Send + 'static)> {
410/// #         unreachable!()
411/// #     }
412/// # }
413/// #
414/// # struct MyCustomJoinHandle;
415/// #
416/// # impl Future for MyCustomJoinHandle {
417/// #     type Output = Result<(), MyCustomJoinError>;
418/// #
419/// #     fn poll(self: Pin<&mut Self>, _cx: &mut Context) -> Poll<Self::Output> {
420/// #         unreachable!()
421/// #     }
422/// # }
423/// #
424/// # struct MyCustomRuntime;
425/// #
426/// # impl MyCustomRuntime {
427/// #     async fn sleep(_: Duration) {
428/// #         unreachable!()
429/// #     }
430/// # }
431/// #
432/// # impl Runtime for MyCustomRuntime {
433/// #     type JoinError = MyCustomJoinError;
434/// #     type JoinHandle = MyCustomJoinHandle;
435/// #
436/// #     fn spawn<F>(fut: F) -> Self::JoinHandle
437/// #     where
438/// #         F: Future<Output = ()> + Send + 'static
439/// #     {
440/// #         unreachable!()
441/// #     }
442/// #
443/// #     fn spawn_blocking<F>(f: F) -> Self::JoinHandle where F: FnOnce() + Send + 'static {
444/// #         unreachable!()
445/// #     }
446/// # }
447/// #
448/// # impl ContextExt for MyCustomRuntime {
449/// #     fn scope<F, R>(locals: TaskLocals, fut: F) -> Pin<Box<dyn Future<Output = R> + Send>>
450/// #     where
451/// #         F: Future<Output = R> + Send + 'static
452/// #     {
453/// #         unreachable!()
454/// #     }
455/// #     fn get_task_locals() -> Option<TaskLocals> {
456/// #         unreachable!()
457/// #     }
458/// # }
459/// #
460/// const PYTHON_CODE: &'static str = r#"
461/// import asyncio
462///
463/// async def py_sleep(duration):
464///     await asyncio.sleep(duration)
465/// "#;
466///
467/// async fn py_sleep(seconds: f32) -> PyResult<()> {
468///     let test_mod = Python::attach(|py| -> PyResult<Py<PyAny>> {
469///         Ok(
470///             PyModule::from_code(
471///                 py,
472///                 &CString::new(PYTHON_CODE).unwrap(),
473///                 &CString::new("test_into_future/test_mod.py").unwrap(),
474///                 &CString::new("test_mod").unwrap(),
475///             )?
476///             .into()
477///         )
478///     })?;
479///
480///     Python::attach(|py| {
481///         pyo3_async_runtimes::generic::into_future::<MyCustomRuntime>(
482///             test_mod
483///                 .call_method1(py, "py_sleep", (seconds,))?
484///                 .into_bound(py),
485///         )
486///     })?
487///     .await?;
488///     Ok(())
489/// }
490/// ```
491pub fn into_future<R>(
492    awaitable: Bound<PyAny>,
493) -> PyResult<impl Future<Output = PyResult<Py<PyAny>>> + Send>
494where
495    R: Runtime + ContextExt,
496{
497    into_future_with_locals(&get_current_locals::<R>(awaitable.py())?, awaitable)
498}
499
500/// Convert a Rust Future into a Python awaitable with a generic runtime
501///
502/// If the `asyncio.Future` returned by this conversion is cancelled via `asyncio.Future.cancel`,
503/// the Rust future will be cancelled as well (new behaviour in `v0.15`).
504///
505/// Python `contextvars` are preserved when calling async Python functions within the Rust future
506/// via [`into_future`] (new behaviour in `v0.15`).
507///
508/// > Although `contextvars` are preserved for async Python functions, synchronous functions will
509/// > unfortunately fail to resolve them when called within the Rust future. This is because the
510/// > function is being called from a Rust thread, not inside an actual Python coroutine context.
511/// >
512/// > As a workaround, you can get the `contextvars` from the current task locals using
513/// > [`get_current_locals`] and [`TaskLocals::context`](`crate::TaskLocals::context`), then wrap your
514/// > synchronous function in a call to `contextvars.Context.run`. This will set the context, call the
515/// > synchronous function, and restore the previous context when it returns or raises an exception.
516///
517/// # Arguments
518/// * `py` - PyO3 GIL guard
519/// * `locals` - The task-local data for Python
520/// * `fut` - The Rust future to be converted
521///
522/// # Examples
523///
524/// ```no_run
525/// # use std::{any::Any, task::{Context, Poll}, pin::Pin, future::Future};
526/// #
527/// # use pyo3_async_runtimes::{
528/// #     TaskLocals,
529/// #     generic::{JoinError, SpawnLocalExt, ContextExt, LocalContextExt, Runtime}
530/// # };
531/// #
532/// # struct MyCustomJoinError;
533/// #
534/// # impl JoinError for MyCustomJoinError {
535/// #     fn is_panic(&self) -> bool {
536/// #         unreachable!()
537/// #     }
538/// #     fn into_panic(self) -> Box<(dyn Any + Send + 'static)> {
539/// #         unreachable!()
540/// #     }
541/// # }
542/// #
543/// # struct MyCustomJoinHandle;
544/// #
545/// # impl Future for MyCustomJoinHandle {
546/// #     type Output = Result<(), MyCustomJoinError>;
547/// #
548/// #     fn poll(self: Pin<&mut Self>, _cx: &mut Context) -> Poll<Self::Output> {
549/// #         unreachable!()
550/// #     }
551/// # }
552/// #
553/// # struct MyCustomRuntime;
554/// #
555/// # impl MyCustomRuntime {
556/// #     async fn sleep(_: Duration) {
557/// #         unreachable!()
558/// #     }
559/// # }
560/// #
561/// # impl Runtime for MyCustomRuntime {
562/// #     type JoinError = MyCustomJoinError;
563/// #     type JoinHandle = MyCustomJoinHandle;
564/// #
565/// #     fn spawn<F>(fut: F) -> Self::JoinHandle
566/// #     where
567/// #         F: Future<Output = ()> + Send + 'static
568/// #     {
569/// #         unreachable!()
570/// #     }
571/// #
572/// #     fn spawn_blocking<F>(f: F) -> Self::JoinHandle where F: FnOnce() + Send + 'static {
573/// #         unreachable!()
574/// #     }
575/// # }
576/// #
577/// # impl ContextExt for MyCustomRuntime {
578/// #     fn scope<F, R>(locals: TaskLocals, fut: F) -> Pin<Box<dyn Future<Output = R> + Send>>
579/// #     where
580/// #         F: Future<Output = R> + Send + 'static
581/// #     {
582/// #         unreachable!()
583/// #     }
584/// #     fn get_task_locals() -> Option<TaskLocals> {
585/// #         unreachable!()
586/// #     }
587/// # }
588/// #
589/// use std::time::Duration;
590///
591/// use pyo3::prelude::*;
592///
593/// /// Awaitable sleep function
594/// #[pyfunction]
595/// fn sleep_for<'p>(py: Python<'p>, secs: Bound<'p, PyAny>) -> PyResult<Bound<'p, PyAny>> {
596///     let secs = secs.extract()?;
597///     pyo3_async_runtimes::generic::future_into_py_with_locals::<MyCustomRuntime, _, _>(
598///         py,
599///         pyo3_async_runtimes::generic::get_current_locals::<MyCustomRuntime>(py)?,
600///         async move {
601///             MyCustomRuntime::sleep(Duration::from_secs(secs)).await;
602///             Ok(())
603///         }
604///     )
605/// }
606/// ```
607#[allow(unused_must_use)]
608pub fn future_into_py_with_locals<R, F, T>(
609    py: Python,
610    locals: TaskLocals,
611    fut: F,
612) -> PyResult<Bound<PyAny>>
613where
614    R: Runtime + ContextExt,
615    F: Future<Output = PyResult<T>> + Send + 'static,
616    T: for<'py> IntoPyObject<'py> + Send + 'static,
617{
618    let (cancel_tx, cancel_rx) = oneshot::channel();
619
620    let py_fut = create_future(locals.0.event_loop.bind(py).clone())?;
621    py_fut.call_method1(
622        pyo3::intern!(py, "add_done_callback"),
623        (PyDoneCallback {
624            cancel_tx: Some(cancel_tx),
625        },),
626    )?;
627
628    let future_tx1: Py<PyAny> = py_fut.clone().into();
629    let future_tx2 = future_tx1.clone_ref(py);
630
631    R::spawn(async move {
632        let locals2 = locals.clone();
633
634        if let Err(e) = R::spawn(async move {
635            let result = R::scope(
636                locals2.clone(),
637                Cancellable::new_with_cancel_rx(fut, cancel_rx),
638            )
639            .await;
640
641            // We should not hold GIL inside async-std/tokio event loop,
642            // because a blocked task may prevent other tasks from progressing.
643            R::spawn_blocking(|| {
644                Python::attach(move |py| {
645                    if cancelled(future_tx1.bind(py))
646                        .map_err(dump_err(py))
647                        .unwrap_or(false)
648                    {
649                        return;
650                    }
651
652                    let _ = set_result(
653                        &locals2.event_loop(py),
654                        future_tx1.bind(py),
655                        result.and_then(|val| val.into_py_any(py)),
656                    )
657                    .map_err(dump_err(py));
658                });
659            });
660        })
661        .await
662        {
663            if e.is_panic() {
664                R::spawn_blocking(|| {
665                    Python::attach(move |py| {
666                        if cancelled(future_tx2.bind(py))
667                            .map_err(dump_err(py))
668                            .unwrap_or(false)
669                        {
670                            return;
671                        }
672
673                        let panic_message = format!(
674                            "rust future panicked: {}",
675                            get_panic_message(&e.into_panic())
676                        );
677                        let _ = set_result(
678                            locals.0.event_loop.bind(py),
679                            future_tx2.bind(py),
680                            Err(RustPanic::new_err(panic_message)),
681                        )
682                        .map_err(dump_err(py));
683                    });
684                });
685            }
686        }
687    });
688
689    Ok(py_fut)
690}
691
692fn get_panic_message(any: &dyn std::any::Any) -> &str {
693    if let Some(str_slice) = any.downcast_ref::<&str>() {
694        str_slice
695    } else if let Some(string) = any.downcast_ref::<String>() {
696        string.as_str()
697    } else {
698        "unknown error"
699    }
700}
701
702pin_project! {
703    /// Future returned by [`timeout`](timeout) and [`timeout_at`](timeout_at).
704    #[must_use = "futures do nothing unless you `.await` or poll them"]
705    #[derive(Debug)]
706    struct Cancellable<T> {
707        #[pin]
708        future: T,
709        #[pin]
710        cancel_rx: oneshot::Receiver<()>,
711
712        poll_cancel_rx: bool
713    }
714}
715
716impl<T> Cancellable<T> {
717    fn new_with_cancel_rx(future: T, cancel_rx: oneshot::Receiver<()>) -> Self {
718        Self {
719            future,
720            cancel_rx,
721
722            poll_cancel_rx: true,
723        }
724    }
725}
726
727impl<'py, F, T> Future for Cancellable<F>
728where
729    F: Future<Output = PyResult<T>>,
730    T: IntoPyObject<'py>,
731{
732    type Output = F::Output;
733
734    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
735        let this = self.project();
736
737        // First, try polling the future
738        if let Poll::Ready(v) = this.future.poll(cx) {
739            return Poll::Ready(v);
740        }
741
742        // Now check for cancellation
743        if *this.poll_cancel_rx {
744            match this.cancel_rx.poll(cx) {
745                Poll::Ready(Ok(())) => {
746                    *this.poll_cancel_rx = false;
747                    // The python future has already been cancelled, so this return value will never
748                    // be used.
749                    Poll::Ready(Err(pyo3::exceptions::PyBaseException::new_err(
750                        "unreachable",
751                    )))
752                }
753                Poll::Ready(Err(_)) => {
754                    *this.poll_cancel_rx = false;
755                    Poll::Pending
756                }
757                Poll::Pending => Poll::Pending,
758            }
759        } else {
760            Poll::Pending
761        }
762    }
763}
764
765#[pyclass]
766struct PyDoneCallback {
767    cancel_tx: Option<oneshot::Sender<()>>,
768}
769
770#[pymethods]
771impl PyDoneCallback {
772    pub fn __call__(&mut self, fut: &Bound<PyAny>) -> PyResult<()> {
773        let py = fut.py();
774
775        if cancelled(fut).map_err(dump_err(py)).unwrap_or(false) {
776            let _ = self.cancel_tx.take().unwrap().send(());
777        }
778
779        Ok(())
780    }
781}
782
783/// Convert a Rust Future into a Python awaitable with a generic runtime
784///
785/// If the `asyncio.Future` returned by this conversion is cancelled via `asyncio.Future.cancel`,
786/// the Rust future will be cancelled as well (new behaviour in `v0.15`).
787///
788/// Python `contextvars` are preserved when calling async Python functions within the Rust future
789/// via [`into_future`] (new behaviour in `v0.15`).
790///
791/// > Although `contextvars` are preserved for async Python functions, synchronous functions will
792/// > unfortunately fail to resolve them when called within the Rust future. This is because the
793/// > function is being called from a Rust thread, not inside an actual Python coroutine context.
794/// >
795/// > As a workaround, you can get the `contextvars` from the current task locals using
796/// > [`get_current_locals`] and [`TaskLocals::context`](`crate::TaskLocals::context`), then wrap your
797/// > synchronous function in a call to `contextvars.Context.run`. This will set the context, call the
798/// > synchronous function, and restore the previous context when it returns or raises an exception.
799///
800/// # Arguments
801/// * `py` - The current PyO3 GIL guard
802/// * `fut` - The Rust future to be converted
803///
804/// # Examples
805///
806/// ```no_run
807/// # use std::{any::Any, task::{Context, Poll}, pin::Pin, future::Future};
808/// #
809/// # use pyo3_async_runtimes::{
810/// #     TaskLocals,
811/// #     generic::{JoinError, SpawnLocalExt, ContextExt, LocalContextExt, Runtime}
812/// # };
813/// #
814/// # struct MyCustomJoinError;
815/// #
816/// # impl JoinError for MyCustomJoinError {
817/// #     fn is_panic(&self) -> bool {
818/// #         unreachable!()
819/// #     }
820/// #     fn into_panic(self) -> Box<(dyn Any + Send + 'static)> {
821/// #         unreachable!()
822/// #     }
823/// # }
824/// #
825/// # struct MyCustomJoinHandle;
826/// #
827/// # impl Future for MyCustomJoinHandle {
828/// #     type Output = Result<(), MyCustomJoinError>;
829/// #
830/// #     fn poll(self: Pin<&mut Self>, _cx: &mut Context) -> Poll<Self::Output> {
831/// #         unreachable!()
832/// #     }
833/// # }
834/// #
835/// # struct MyCustomRuntime;
836/// #
837/// # impl MyCustomRuntime {
838/// #     async fn sleep(_: Duration) {
839/// #         unreachable!()
840/// #     }
841/// # }
842/// #
843/// # impl Runtime for MyCustomRuntime {
844/// #     type JoinError = MyCustomJoinError;
845/// #     type JoinHandle = MyCustomJoinHandle;
846/// #
847/// #     fn spawn<F>(fut: F) -> Self::JoinHandle
848/// #     where
849/// #         F: Future<Output = ()> + Send + 'static
850/// #     {
851/// #         unreachable!()
852/// #     }
853/// #
854/// #     fn spawn_blocking<F>(f: F) -> Self::JoinHandle where F: FnOnce() + Send + 'static {
855/// #         unreachable!()
856/// #     }
857/// # }
858/// #
859/// # impl ContextExt for MyCustomRuntime {
860/// #     fn scope<F, R>(locals: TaskLocals, fut: F) -> Pin<Box<dyn Future<Output = R> + Send>>
861/// #     where
862/// #         F: Future<Output = R> + Send + 'static
863/// #     {
864/// #         unreachable!()
865/// #     }
866/// #     fn get_task_locals() -> Option<TaskLocals> {
867/// #         unreachable!()
868/// #     }
869/// # }
870/// #
871/// use std::time::Duration;
872///
873/// use pyo3::prelude::*;
874///
875/// /// Awaitable sleep function
876/// #[pyfunction]
877/// fn sleep_for<'p>(py: Python<'p>, secs: Bound<'p, PyAny>) -> PyResult<Bound<'p, PyAny>> {
878///     let secs = secs.extract()?;
879///     pyo3_async_runtimes::generic::future_into_py::<MyCustomRuntime, _, _>(py, async move {
880///         MyCustomRuntime::sleep(Duration::from_secs(secs)).await;
881///         Ok(())
882///     })
883/// }
884/// ```
885pub fn future_into_py<R, F, T>(py: Python, fut: F) -> PyResult<Bound<PyAny>>
886where
887    R: Runtime + ContextExt,
888    F: Future<Output = PyResult<T>> + Send + 'static,
889    T: for<'py> IntoPyObject<'py> + Send + 'static,
890{
891    future_into_py_with_locals::<R, F, T>(py, get_current_locals::<R>(py)?, fut)
892}
893
894/// Convert a `!Send` Rust Future into a Python awaitable with a generic runtime and manual
895/// specification of task locals.
896///
897/// If the `asyncio.Future` returned by this conversion is cancelled via `asyncio.Future.cancel`,
898/// the Rust future will be cancelled as well (new behaviour in `v0.15`).
899///
900/// Python `contextvars` are preserved when calling async Python functions within the Rust future
901/// via [`into_future`] (new behaviour in `v0.15`).
902///
903/// > Although `contextvars` are preserved for async Python functions, synchronous functions will
904/// > unfortunately fail to resolve them when called within the Rust future. This is because the
905/// > function is being called from a Rust thread, not inside an actual Python coroutine context.
906/// >
907/// > As a workaround, you can get the `contextvars` from the current task locals using
908/// > [`get_current_locals`] and [`TaskLocals::context`](`crate::TaskLocals::context`), then wrap your
909/// > synchronous function in a call to `contextvars.Context.run`. This will set the context, call the
910/// > synchronous function, and restore the previous context when it returns or raises an exception.
911///
912/// # Arguments
913/// * `py` - PyO3 GIL guard
914/// * `locals` - The task locals for the future
915/// * `fut` - The Rust future to be converted
916///
917/// # Examples
918///
919/// ```no_run
920/// # use std::{any::Any, task::{Context, Poll}, pin::Pin, future::Future};
921/// #
922/// # use pyo3_async_runtimes::{
923/// #     TaskLocals,
924/// #     generic::{JoinError, SpawnLocalExt, ContextExt, LocalContextExt, Runtime}
925/// # };
926/// #
927/// # struct MyCustomJoinError;
928/// #
929/// # impl JoinError for MyCustomJoinError {
930/// #     fn is_panic(&self) -> bool {
931/// #         unreachable!()
932/// #     }
933/// #     fn into_panic(self) -> Box<(dyn Any + Send + 'static)> {
934/// #         unreachable!()
935/// #     }
936/// # }
937/// #
938/// # struct MyCustomJoinHandle;
939/// #
940/// # impl Future for MyCustomJoinHandle {
941/// #     type Output = Result<(), MyCustomJoinError>;
942/// #
943/// #     fn poll(self: Pin<&mut Self>, _cx: &mut Context) -> Poll<Self::Output> {
944/// #         unreachable!()
945/// #     }
946/// # }
947/// #
948/// # struct MyCustomRuntime;
949/// #
950/// # impl MyCustomRuntime {
951/// #     async fn sleep(_: Duration) {
952/// #         unreachable!()
953/// #     }
954/// # }
955/// #
956/// # impl Runtime for MyCustomRuntime {
957/// #     type JoinError = MyCustomJoinError;
958/// #     type JoinHandle = MyCustomJoinHandle;
959/// #
960/// #     fn spawn<F>(fut: F) -> Self::JoinHandle
961/// #     where
962/// #         F: Future<Output = ()> + Send + 'static
963/// #     {
964/// #         unreachable!()
965/// #     }
966/// #
967/// #     fn spawn_blocking<F>(f: F) -> Self::JoinHandle where F: FnOnce() + Send + 'static {
968/// #         unreachable!()
969/// #     }
970/// # }
971/// #
972/// # impl ContextExt for MyCustomRuntime {
973/// #     fn scope<F, R>(locals: TaskLocals, fut: F) -> Pin<Box<dyn Future<Output = R> + Send>>
974/// #     where
975/// #         F: Future<Output = R> + Send + 'static
976/// #     {
977/// #         unreachable!()
978/// #     }
979/// #     fn get_task_locals() -> Option<TaskLocals> {
980/// #         unreachable!()
981/// #     }
982/// # }
983/// #
984/// # impl SpawnLocalExt for MyCustomRuntime {
985/// #     fn spawn_local<F>(fut: F) -> Self::JoinHandle
986/// #     where
987/// #         F: Future<Output = ()> + 'static
988/// #     {
989/// #         unreachable!()
990/// #     }
991/// # }
992/// #
993/// # impl LocalContextExt for MyCustomRuntime {
994/// #     fn scope_local<F, R>(locals: TaskLocals, fut: F) -> Pin<Box<dyn Future<Output = R>>>
995/// #     where
996/// #         F: Future<Output = R> + 'static
997/// #     {
998/// #         unreachable!()
999/// #     }
1000/// # }
1001/// #
1002/// use std::{rc::Rc, time::Duration};
1003///
1004/// use pyo3::prelude::*;
1005///
1006/// /// Awaitable sleep function
1007/// #[pyfunction]
1008/// fn sleep_for(py: Python, secs: u64) -> PyResult<Bound<PyAny>> {
1009///     // Rc is !Send so it cannot be passed into pyo3_async_runtimes::generic::future_into_py
1010///     let secs = Rc::new(secs);
1011///
1012///     pyo3_async_runtimes::generic::local_future_into_py_with_locals::<MyCustomRuntime, _, _>(
1013///         py,
1014///         pyo3_async_runtimes::generic::get_current_locals::<MyCustomRuntime>(py)?,
1015///         async move {
1016///             MyCustomRuntime::sleep(Duration::from_secs(*secs)).await;
1017///             Ok(())
1018///         }
1019///     )
1020/// }
1021/// ```
1022#[deprecated(
1023    since = "0.18.0",
1024    note = "Questionable whether these conversions have real-world utility (see https://github.com/awestlake87/pyo3-asyncio/issues/59#issuecomment-1008038497 and let me know if you disagree!)"
1025)]
1026#[allow(unused_must_use)]
1027pub fn local_future_into_py_with_locals<R, F, T>(
1028    py: Python,
1029    locals: TaskLocals,
1030    fut: F,
1031) -> PyResult<Bound<PyAny>>
1032where
1033    R: Runtime + SpawnLocalExt + LocalContextExt,
1034    F: Future<Output = PyResult<T>> + 'static,
1035    T: for<'py> IntoPyObject<'py>,
1036{
1037    let (cancel_tx, cancel_rx) = oneshot::channel();
1038
1039    let py_fut = create_future(locals.0.event_loop.clone_ref(py).into_bound(py))?;
1040    py_fut.call_method1(
1041        pyo3::intern!(py, "add_done_callback"),
1042        (PyDoneCallback {
1043            cancel_tx: Some(cancel_tx),
1044        },),
1045    )?;
1046
1047    let future_tx1: Py<PyAny> = py_fut.clone().into();
1048    let future_tx2 = future_tx1.clone_ref(py);
1049
1050    R::spawn_local(async move {
1051        let locals2 = locals.clone();
1052
1053        if let Err(e) = R::spawn_local(async move {
1054            let result = R::scope_local(
1055                locals2.clone(),
1056                Cancellable::new_with_cancel_rx(fut, cancel_rx),
1057            )
1058            .await;
1059
1060            Python::attach(move |py| {
1061                if cancelled(future_tx1.bind(py))
1062                    .map_err(dump_err(py))
1063                    .unwrap_or(false)
1064                {
1065                    return;
1066                }
1067
1068                let _ = set_result(
1069                    locals2.0.event_loop.bind(py),
1070                    future_tx1.bind(py),
1071                    result.and_then(|val| val.into_py_any(py)),
1072                )
1073                .map_err(dump_err(py));
1074            });
1075        })
1076        .await
1077        {
1078            if e.is_panic() {
1079                Python::attach(move |py| {
1080                    if cancelled(future_tx2.bind(py))
1081                        .map_err(dump_err(py))
1082                        .unwrap_or(false)
1083                    {
1084                        return;
1085                    }
1086
1087                    let panic_message = format!(
1088                        "rust future panicked: {}",
1089                        get_panic_message(&e.into_panic())
1090                    );
1091                    let _ = set_result(
1092                        locals.0.event_loop.bind(py),
1093                        future_tx2.bind(py),
1094                        Err(RustPanic::new_err(panic_message)),
1095                    )
1096                    .map_err(dump_err(py));
1097                });
1098            }
1099        }
1100    });
1101
1102    Ok(py_fut)
1103}
1104
1105/// Convert a `!Send` Rust Future into a Python awaitable with a generic runtime
1106///
1107/// If the `asyncio.Future` returned by this conversion is cancelled via `asyncio.Future.cancel`,
1108/// the Rust future will be cancelled as well (new behaviour in `v0.15`).
1109///
1110/// Python `contextvars` are preserved when calling async Python functions within the Rust future
1111/// via [`into_future`] (new behaviour in `v0.15`).
1112///
1113/// > Although `contextvars` are preserved for async Python functions, synchronous functions will
1114/// > unfortunately fail to resolve them when called within the Rust future. This is because the
1115/// > function is being called from a Rust thread, not inside an actual Python coroutine context.
1116/// >
1117/// > As a workaround, you can get the `contextvars` from the current task locals using
1118/// > [`get_current_locals`] and [`TaskLocals::context`](`crate::TaskLocals::context`), then wrap your
1119/// > synchronous function in a call to `contextvars.Context.run`. This will set the context, call the
1120/// > synchronous function, and restore the previous context when it returns or raises an exception.
1121///
1122/// # Arguments
1123/// * `py` - The current PyO3 GIL guard
1124/// * `fut` - The Rust future to be converted
1125///
1126/// # Examples
1127///
1128/// ```no_run
1129/// # use std::{any::Any, task::{Context, Poll}, pin::Pin, future::Future};
1130/// #
1131/// # use pyo3_async_runtimes::{
1132/// #     TaskLocals,
1133/// #     generic::{JoinError, SpawnLocalExt, ContextExt, LocalContextExt, Runtime}
1134/// # };
1135/// #
1136/// # struct MyCustomJoinError;
1137/// #
1138/// # impl JoinError for MyCustomJoinError {
1139/// #     fn is_panic(&self) -> bool {
1140/// #         unreachable!()
1141/// #     }
1142/// #     fn into_panic(self) -> Box<(dyn Any + Send + 'static)> {
1143/// #         unreachable!()
1144/// #     }
1145/// # }
1146/// #
1147/// # struct MyCustomJoinHandle;
1148/// #
1149/// # impl Future for MyCustomJoinHandle {
1150/// #     type Output = Result<(), MyCustomJoinError>;
1151/// #
1152/// #     fn poll(self: Pin<&mut Self>, _cx: &mut Context) -> Poll<Self::Output> {
1153/// #         unreachable!()
1154/// #     }
1155/// # }
1156/// #
1157/// # struct MyCustomRuntime;
1158/// #
1159/// # impl MyCustomRuntime {
1160/// #     async fn sleep(_: Duration) {
1161/// #         unreachable!()
1162/// #     }
1163/// # }
1164/// #
1165/// # impl Runtime for MyCustomRuntime {
1166/// #     type JoinError = MyCustomJoinError;
1167/// #     type JoinHandle = MyCustomJoinHandle;
1168/// #
1169/// #     fn spawn<F>(fut: F) -> Self::JoinHandle
1170/// #     where
1171/// #         F: Future<Output = ()> + Send + 'static
1172/// #     {
1173/// #         unreachable!()
1174/// #     }
1175/// #
1176/// #     fn spawn_blocking<F>(f: F) -> Self::JoinHandle where F: FnOnce() + Send + 'static {
1177/// #         unreachable!()
1178/// #     }
1179/// # }
1180/// #
1181/// # impl ContextExt for MyCustomRuntime {
1182/// #     fn scope<F, R>(locals: TaskLocals, fut: F) -> Pin<Box<dyn Future<Output = R> + Send>>
1183/// #     where
1184/// #         F: Future<Output = R> + Send + 'static
1185/// #     {
1186/// #         unreachable!()
1187/// #     }
1188/// #     fn get_task_locals() -> Option<TaskLocals> {
1189/// #         unreachable!()
1190/// #     }
1191/// # }
1192/// #
1193/// # impl SpawnLocalExt for MyCustomRuntime {
1194/// #     fn spawn_local<F>(fut: F) -> Self::JoinHandle
1195/// #     where
1196/// #         F: Future<Output = ()> + 'static
1197/// #     {
1198/// #         unreachable!()
1199/// #     }
1200/// # }
1201/// #
1202/// # impl LocalContextExt for MyCustomRuntime {
1203/// #     fn scope_local<F, R>(locals: TaskLocals, fut: F) -> Pin<Box<dyn Future<Output = R>>>
1204/// #     where
1205/// #         F: Future<Output = R> + 'static
1206/// #     {
1207/// #         unreachable!()
1208/// #     }
1209/// # }
1210/// #
1211/// use std::{rc::Rc, time::Duration};
1212///
1213/// use pyo3::prelude::*;
1214///
1215/// /// Awaitable sleep function
1216/// #[pyfunction]
1217/// fn sleep_for(py: Python, secs: u64) -> PyResult<Bound<PyAny>> {
1218///     // Rc is !Send so it cannot be passed into pyo3_async_runtimes::generic::future_into_py
1219///     let secs = Rc::new(secs);
1220///
1221///     pyo3_async_runtimes::generic::local_future_into_py::<MyCustomRuntime, _, _>(py, async move {
1222///         MyCustomRuntime::sleep(Duration::from_secs(*secs)).await;
1223///         Ok(())
1224///     })
1225/// }
1226/// ```
1227#[deprecated(
1228    since = "0.18.0",
1229    note = "Questionable whether these conversions have real-world utility (see https://github.com/awestlake87/pyo3-asyncio/issues/59#issuecomment-1008038497 and let me know if you disagree!)"
1230)]
1231#[allow(deprecated)]
1232pub fn local_future_into_py<R, F, T>(py: Python, fut: F) -> PyResult<Bound<PyAny>>
1233where
1234    R: Runtime + ContextExt + SpawnLocalExt + LocalContextExt,
1235    F: Future<Output = PyResult<T>> + 'static,
1236    T: for<'py> IntoPyObject<'py>,
1237{
1238    local_future_into_py_with_locals::<R, F, T>(py, get_current_locals::<R>(py)?, fut)
1239}
1240
1241/// <span class="module-item stab portability" style="display: inline; border-radius: 3px; padding: 2px; font-size: 80%; line-height: 1.2;"><code>unstable-streams</code></span> Convert an async generator into a stream
1242///
1243/// **This API is marked as unstable** and is only available when the
1244/// `unstable-streams` crate feature is enabled. This comes with no
1245/// stability guarantees, and could be changed or removed at any time.
1246///
1247/// # Arguments
1248/// * `locals` - The current task locals
1249/// * `gen` - The Python async generator to be converted
1250///
1251/// # Examples
1252/// ```no_run
1253/// # use std::{any::Any, task::{Context, Poll}, pin::Pin, future::Future};
1254/// #
1255/// # use pyo3_async_runtimes::{
1256/// #     TaskLocals,
1257/// #     generic::{JoinError, ContextExt, Runtime}
1258/// # };
1259/// #
1260/// # struct MyCustomJoinError;
1261/// #
1262/// # impl JoinError for MyCustomJoinError {
1263/// #     fn is_panic(&self) -> bool {
1264/// #         unreachable!()
1265/// #     }
1266/// #     fn into_panic(self) -> Box<(dyn Any + Send + 'static)> {
1267/// #         unreachable!()
1268/// #     }
1269/// # }
1270/// #
1271/// # struct MyCustomJoinHandle;
1272/// #
1273/// # impl Future for MyCustomJoinHandle {
1274/// #     type Output = Result<(), MyCustomJoinError>;
1275/// #
1276/// #     fn poll(self: Pin<&mut Self>, _cx: &mut Context) -> Poll<Self::Output> {
1277/// #         unreachable!()
1278/// #     }
1279/// # }
1280/// #
1281/// # struct MyCustomRuntime;
1282/// #
1283/// # impl Runtime for MyCustomRuntime {
1284/// #     type JoinError = MyCustomJoinError;
1285/// #     type JoinHandle = MyCustomJoinHandle;
1286/// #
1287/// #     fn spawn<F>(fut: F) -> Self::JoinHandle
1288/// #     where
1289/// #         F: Future<Output = ()> + Send + 'static
1290/// #     {
1291/// #         unreachable!()
1292/// #     }
1293/// #
1294/// #     fn spawn_blocking<F>(f: F) -> Self::JoinHandle where F: FnOnce() + Send + 'static {
1295/// #         unreachable!()
1296/// #     }
1297/// # }
1298/// #
1299/// # impl ContextExt for MyCustomRuntime {
1300/// #     fn scope<F, R>(locals: TaskLocals, fut: F) -> Pin<Box<dyn Future<Output = R> + Send>>
1301/// #     where
1302/// #         F: Future<Output = R> + Send + 'static
1303/// #     {
1304/// #         unreachable!()
1305/// #     }
1306/// #     fn get_task_locals() -> Option<TaskLocals> {
1307/// #         unreachable!()
1308/// #     }
1309/// # }
1310///
1311/// use pyo3::prelude::*;
1312/// use futures_util::stream::{StreamExt, TryStreamExt};
1313/// use std::ffi::CString;
1314///
1315/// const TEST_MOD: &str = r#"
1316/// import asyncio
1317///
1318/// async def gen():
1319///     for i in range(10):
1320///         await asyncio.sleep(0.1)
1321///         yield i
1322/// "#;
1323///
1324/// # async fn test_async_gen() -> PyResult<()> {
1325/// let stream = Python::attach(|py| {
1326///     let test_mod = PyModule::from_code(
1327///         py,
1328///         &CString::new(TEST_MOD).unwrap(),
1329///         &CString::new("test_rust_coroutine/test_mod.py").unwrap(),
1330///         &CString::new("test_mod").unwrap(),
1331///     )?;
1332///
1333///     pyo3_async_runtimes::generic::into_stream_with_locals_v1::<MyCustomRuntime>(
1334///         pyo3_async_runtimes::generic::get_current_locals::<MyCustomRuntime>(py)?,
1335///         test_mod.call_method0("gen")?
1336///     )
1337/// })?;
1338///
1339/// let vals = stream
1340///     .map(|item| Python::attach(|py| -> PyResult<i32> { Ok(item?.bind(py).extract()?) }))
1341///     .try_collect::<Vec<i32>>()
1342///     .await?;
1343///
1344/// assert_eq!((0..10).collect::<Vec<i32>>(), vals);
1345///
1346/// Ok(())
1347/// # }
1348/// ```
1349#[cfg(feature = "unstable-streams")]
1350#[allow(unused_must_use)] // False positive unused lint on `R::spawn`
1351pub fn into_stream_with_locals_v1<R>(
1352    locals: TaskLocals,
1353    gen: Bound<'_, PyAny>,
1354) -> PyResult<impl futures_util::stream::Stream<Item = PyResult<Py<PyAny>>> + 'static>
1355where
1356    R: Runtime,
1357{
1358    let (tx, rx) = async_channel::bounded(1);
1359    let py = gen.py();
1360    let anext: Py<PyAny> = gen.getattr(pyo3::intern!(py, "__anext__"))?.into();
1361
1362    R::spawn(async move {
1363        loop {
1364            let fut = Python::attach(|py| -> PyResult<_> {
1365                into_future_with_locals(&locals, anext.bind(py).call0()?)
1366            });
1367            let item = match fut {
1368                Ok(fut) => match fut.await {
1369                    Ok(item) => Ok(item),
1370                    Err(e) => {
1371                        let stop_iter = Python::attach(|py| {
1372                            e.is_instance_of::<pyo3::exceptions::PyStopAsyncIteration>(py)
1373                        });
1374
1375                        if stop_iter {
1376                            // end the iteration
1377                            break;
1378                        } else {
1379                            Err(e)
1380                        }
1381                    }
1382                },
1383                Err(e) => Err(e),
1384            };
1385
1386            if tx.send(item).await.is_err() {
1387                // receiving side was dropped
1388                break;
1389            }
1390        }
1391    });
1392
1393    Ok(rx)
1394}
1395
1396/// <span class="module-item stab portability" style="display: inline; border-radius: 3px; padding: 2px; font-size: 80%; line-height: 1.2;"><code>unstable-streams</code></span> Convert an async generator into a stream
1397///
1398/// **This API is marked as unstable** and is only available when the
1399/// `unstable-streams` crate feature is enabled. This comes with no
1400/// stability guarantees, and could be changed or removed at any time.
1401///
1402/// # Arguments
1403/// * `gen` - The Python async generator to be converted
1404///
1405/// # Examples
1406/// ```no_run
1407/// # use std::{any::Any, task::{Context, Poll}, pin::Pin, future::Future};
1408/// #
1409/// # use pyo3_async_runtimes::{
1410/// #     TaskLocals,
1411/// #     generic::{JoinError, ContextExt, Runtime}
1412/// # };
1413/// #
1414/// # struct MyCustomJoinError;
1415/// #
1416/// # impl JoinError for MyCustomJoinError {
1417/// #     fn is_panic(&self) -> bool {
1418/// #         unreachable!()
1419/// #     }
1420/// #     fn into_panic(self) -> Box<(dyn Any + Send + 'static)> {
1421/// #         unreachable!()
1422/// #     }
1423/// # }
1424/// #
1425/// # struct MyCustomJoinHandle;
1426/// #
1427/// # impl Future for MyCustomJoinHandle {
1428/// #     type Output = Result<(), MyCustomJoinError>;
1429/// #
1430/// #     fn poll(self: Pin<&mut Self>, _cx: &mut Context) -> Poll<Self::Output> {
1431/// #         unreachable!()
1432/// #     }
1433/// # }
1434/// #
1435/// # struct MyCustomRuntime;
1436/// #
1437/// # impl Runtime for MyCustomRuntime {
1438/// #     type JoinError = MyCustomJoinError;
1439/// #     type JoinHandle = MyCustomJoinHandle;
1440/// #
1441/// #     fn spawn<F>(fut: F) -> Self::JoinHandle
1442/// #     where
1443/// #         F: Future<Output = ()> + Send + 'static
1444/// #     {
1445/// #         unreachable!()
1446/// #     }
1447/// #
1448/// #     fn spawn_blocking<F>(f: F) -> Self::JoinHandle where F: FnOnce() + Send + 'static {
1449/// #         unreachable!()
1450/// #     }
1451/// # }
1452/// #
1453/// # impl ContextExt for MyCustomRuntime {
1454/// #     fn scope<F, R>(locals: TaskLocals, fut: F) -> Pin<Box<dyn Future<Output = R> + Send>>
1455/// #     where
1456/// #         F: Future<Output = R> + Send + 'static
1457/// #     {
1458/// #         unreachable!()
1459/// #     }
1460/// #     fn get_task_locals() -> Option<TaskLocals> {
1461/// #         unreachable!()
1462/// #     }
1463/// # }
1464///
1465/// use pyo3::prelude::*;
1466/// use futures_util::stream::{StreamExt, TryStreamExt};
1467/// use std::ffi::CString;
1468///
1469/// const TEST_MOD: &str = r#"
1470/// import asyncio
1471///
1472/// async def gen():
1473///     for i in range(10):
1474///         await asyncio.sleep(0.1)
1475///         yield i
1476/// "#;
1477///
1478/// # async fn test_async_gen() -> PyResult<()> {
1479/// let stream = Python::attach(|py| {
1480///     let test_mod = PyModule::from_code(
1481///         py,
1482///         &CString::new(TEST_MOD).unwrap(),
1483///         &CString::new("test_rust_coroutine/test_mod.py").unwrap(),
1484///         &CString::new("test_mod").unwrap(),
1485///     )?;
1486///
1487///     pyo3_async_runtimes::generic::into_stream_v1::<MyCustomRuntime>(test_mod.call_method0("gen")?)
1488/// })?;
1489///
1490/// let vals = stream
1491///     .map(|item| Python::attach(|py| -> PyResult<i32> { Ok(item?.bind(py).extract()?) }))
1492///     .try_collect::<Vec<i32>>()
1493///     .await?;
1494///
1495/// assert_eq!((0..10).collect::<Vec<i32>>(), vals);
1496///
1497/// Ok(())
1498/// # }
1499/// ```
1500#[cfg(feature = "unstable-streams")]
1501pub fn into_stream_v1<R>(
1502    gen: Bound<'_, PyAny>,
1503) -> PyResult<impl futures_util::stream::Stream<Item = PyResult<Py<PyAny>>> + 'static>
1504where
1505    R: Runtime + ContextExt,
1506{
1507    into_stream_with_locals_v1::<R>(get_current_locals::<R>(gen.py())?, gen)
1508}
1509
1510trait Sender: Send + 'static {
1511    fn send(&mut self, py: Python, locals: TaskLocals, item: Py<PyAny>) -> PyResult<Py<PyAny>>;
1512    fn close(&mut self) -> PyResult<()>;
1513}
1514
1515#[cfg(feature = "unstable-streams")]
1516struct GenericSender<R>
1517where
1518    R: Runtime,
1519{
1520    runtime: PhantomData<R>,
1521    tx: mpsc::Sender<Py<PyAny>>,
1522}
1523
1524#[cfg(feature = "unstable-streams")]
1525impl<R> Sender for GenericSender<R>
1526where
1527    R: Runtime + ContextExt,
1528{
1529    fn send(&mut self, py: Python, locals: TaskLocals, item: Py<PyAny>) -> PyResult<Py<PyAny>> {
1530        match self.tx.try_send(item.clone_ref(py)) {
1531            Ok(_) => true.into_py_any(py),
1532            Err(e) => {
1533                if e.is_full() {
1534                    let mut tx = self.tx.clone();
1535
1536                    future_into_py_with_locals::<R, _, bool>(py, locals, async move {
1537                        if tx.flush().await.is_err() {
1538                            // receiving side disconnected
1539                            return Ok(false);
1540                        }
1541                        if tx.send(item).await.is_err() {
1542                            // receiving side disconnected
1543                            return Ok(false);
1544                        }
1545                        Ok(true)
1546                    })
1547                    .map(Bound::unbind)
1548                } else {
1549                    false.into_py_any(py)
1550                }
1551            }
1552        }
1553    }
1554    fn close(&mut self) -> PyResult<()> {
1555        self.tx.close_channel();
1556        Ok(())
1557    }
1558}
1559
1560#[pyclass]
1561struct SenderGlue {
1562    locals: TaskLocals,
1563    tx: Arc<Mutex<dyn Sender>>,
1564}
1565#[pymethods]
1566impl SenderGlue {
1567    pub fn send(&mut self, item: Py<PyAny>) -> PyResult<Py<PyAny>> {
1568        Python::attach(|py| self.tx.lock().unwrap().send(py, self.locals.clone(), item))
1569    }
1570    pub fn close(&mut self) -> PyResult<()> {
1571        self.tx.lock().unwrap().close()
1572    }
1573}
1574
1575#[cfg(feature = "unstable-streams")]
1576const STREAM_GLUE: &str = r#"
1577import inspect
1578
1579async def forward(gen, sender):
1580    async for item in gen:
1581        should_continue = sender.send(item)
1582
1583        if inspect.isawaitable(should_continue):
1584            should_continue = await should_continue
1585
1586        if should_continue:
1587            continue
1588        else:
1589            break
1590
1591    sender.close()
1592"#;
1593
1594/// <span class="module-item stab portability" style="display: inline; border-radius: 3px; padding: 2px; font-size: 80%; line-height: 1.2;"><code>unstable-streams</code></span> Convert an async generator into a stream
1595///
1596/// **This API is marked as unstable** and is only available when the
1597/// `unstable-streams` crate feature is enabled. This comes with no
1598/// stability guarantees, and could be changed or removed at any time.
1599///
1600/// # Arguments
1601/// * `locals` - The current task locals
1602/// * `gen` - The Python async generator to be converted
1603///
1604/// # Examples
1605/// ```no_run
1606/// # use std::{any::Any, task::{Context, Poll}, pin::Pin, future::Future};
1607/// #
1608/// # use pyo3_async_runtimes::{
1609/// #     TaskLocals,
1610/// #     generic::{JoinError, ContextExt, Runtime}
1611/// # };
1612/// #
1613/// # struct MyCustomJoinError;
1614/// #
1615/// # impl JoinError for MyCustomJoinError {
1616/// #     fn is_panic(&self) -> bool {
1617/// #         unreachable!()
1618/// #     }
1619/// #     fn into_panic(self) -> Box<(dyn Any + Send + 'static)> {
1620/// #         unreachable!()
1621/// #     }
1622/// # }
1623/// #
1624/// # struct MyCustomJoinHandle;
1625/// #
1626/// # impl Future for MyCustomJoinHandle {
1627/// #     type Output = Result<(), MyCustomJoinError>;
1628/// #
1629/// #     fn poll(self: Pin<&mut Self>, _cx: &mut Context) -> Poll<Self::Output> {
1630/// #         unreachable!()
1631/// #     }
1632/// # }
1633/// #
1634/// # struct MyCustomRuntime;
1635/// #
1636/// # impl Runtime for MyCustomRuntime {
1637/// #     type JoinError = MyCustomJoinError;
1638/// #     type JoinHandle = MyCustomJoinHandle;
1639/// #
1640/// #     fn spawn<F>(fut: F) -> Self::JoinHandle
1641/// #     where
1642/// #         F: Future<Output = ()> + Send + 'static
1643/// #     {
1644/// #         unreachable!()
1645/// #     }
1646/// #
1647/// #     fn spawn_blocking<F>(f: F) -> Self::JoinHandle where F: FnOnce() + Send + 'static {
1648/// #         unreachable!()
1649/// #     }
1650/// # }
1651/// #
1652/// # impl ContextExt for MyCustomRuntime {
1653/// #     fn scope<F, R>(locals: TaskLocals, fut: F) -> Pin<Box<dyn Future<Output = R> + Send>>
1654/// #     where
1655/// #         F: Future<Output = R> + Send + 'static
1656/// #     {
1657/// #         unreachable!()
1658/// #     }
1659/// #     fn get_task_locals() -> Option<TaskLocals> {
1660/// #         unreachable!()
1661/// #     }
1662/// # }
1663///
1664/// use pyo3::prelude::*;
1665/// use futures_util::stream::{StreamExt, TryStreamExt};
1666/// use std::ffi::CString;
1667///
1668/// const TEST_MOD: &str = r#"
1669/// import asyncio
1670///
1671/// async def gen():
1672///     for i in range(10):
1673///         await asyncio.sleep(0.1)
1674///         yield i
1675/// "#;
1676///
1677/// # async fn test_async_gen() -> PyResult<()> {
1678/// let stream = Python::attach(|py| {
1679///     let test_mod = PyModule::from_code(
1680///         py,
1681///         &CString::new(TEST_MOD).unwrap(),
1682///         &CString::new("test_rust_coroutine/test_mod.py").unwrap(),
1683///         &CString::new("test_mod").unwrap(),
1684///     )?;
1685///
1686///     pyo3_async_runtimes::generic::into_stream_with_locals_v2::<MyCustomRuntime>(
1687///         pyo3_async_runtimes::generic::get_current_locals::<MyCustomRuntime>(py)?,
1688///         test_mod.call_method0("gen")?
1689///     )
1690/// })?;
1691///
1692/// let vals = stream
1693///     .map(|item| Python::attach(|py| -> PyResult<i32> { Ok(item.bind(py).extract()?) }))
1694///     .try_collect::<Vec<i32>>()
1695///     .await?;
1696///
1697/// assert_eq!((0..10).collect::<Vec<i32>>(), vals);
1698///
1699/// Ok(())
1700/// # }
1701/// ```
1702#[cfg(feature = "unstable-streams")]
1703pub fn into_stream_with_locals_v2<R>(
1704    locals: TaskLocals,
1705    gen: Bound<'_, PyAny>,
1706) -> PyResult<impl futures_util::stream::Stream<Item = Py<PyAny>> + 'static>
1707where
1708    R: Runtime + ContextExt,
1709{
1710    use std::ffi::CString;
1711
1712    use pyo3::sync::PyOnceLock;
1713
1714    static GLUE_MOD: PyOnceLock<Py<PyAny>> = PyOnceLock::new();
1715    let py = gen.py();
1716    let glue = GLUE_MOD
1717        .get_or_try_init(py, || -> PyResult<Py<PyAny>> {
1718            Ok(PyModule::from_code(
1719                py,
1720                &CString::new(STREAM_GLUE).unwrap(),
1721                &CString::new("pyo3_async_runtimes/pyo3_async_runtimes_glue.py").unwrap(),
1722                &CString::new("pyo3_async_runtimes_glue").unwrap(),
1723            )?
1724            .into())
1725        })?
1726        .bind(py);
1727
1728    let (tx, rx) = mpsc::channel(10);
1729
1730    locals.event_loop(py).call_method1(
1731        pyo3::intern!(py, "call_soon_threadsafe"),
1732        (
1733            locals
1734                .event_loop(py)
1735                .getattr(pyo3::intern!(py, "create_task"))?,
1736            glue.call_method1(
1737                pyo3::intern!(py, "forward"),
1738                (
1739                    gen,
1740                    SenderGlue {
1741                        locals,
1742                        tx: Arc::new(Mutex::new(GenericSender {
1743                            runtime: PhantomData::<R>,
1744                            tx,
1745                        })),
1746                    },
1747                ),
1748            )?,
1749        ),
1750    )?;
1751    Ok(rx)
1752}
1753
1754/// <span class="module-item stab portability" style="display: inline; border-radius: 3px; padding: 2px; font-size: 80%; line-height: 1.2;"><code>unstable-streams</code></span> Convert an async generator into a stream
1755///
1756/// **This API is marked as unstable** and is only available when the
1757/// `unstable-streams` crate feature is enabled. This comes with no
1758/// stability guarantees, and could be changed or removed at any time.
1759///
1760/// # Arguments
1761/// * `gen` - The Python async generator to be converted
1762///
1763/// # Examples
1764/// ```no_run
1765/// # use std::{any::Any, task::{Context, Poll}, pin::Pin, future::Future};
1766/// #
1767/// # use pyo3_async_runtimes::{
1768/// #     TaskLocals,
1769/// #     generic::{JoinError, ContextExt, Runtime}
1770/// # };
1771/// #
1772/// # struct MyCustomJoinError;
1773/// #
1774/// # impl JoinError for MyCustomJoinError {
1775/// #     fn is_panic(&self) -> bool {
1776/// #         unreachable!()
1777/// #     }
1778/// #     fn into_panic(self) -> Box<(dyn Any + Send + 'static)> {
1779/// #         unreachable!()
1780/// #     }
1781/// # }
1782/// #
1783/// # struct MyCustomJoinHandle;
1784/// #
1785/// # impl Future for MyCustomJoinHandle {
1786/// #     type Output = Result<(), MyCustomJoinError>;
1787/// #
1788/// #     fn poll(self: Pin<&mut Self>, _cx: &mut Context) -> Poll<Self::Output> {
1789/// #         unreachable!()
1790/// #     }
1791/// # }
1792/// #
1793/// # struct MyCustomRuntime;
1794/// #
1795/// # impl Runtime for MyCustomRuntime {
1796/// #     type JoinError = MyCustomJoinError;
1797/// #     type JoinHandle = MyCustomJoinHandle;
1798/// #
1799/// #     fn spawn<F>(fut: F) -> Self::JoinHandle
1800/// #     where
1801/// #         F: Future<Output = ()> + Send + 'static
1802/// #     {
1803/// #         unreachable!()
1804/// #     }
1805/// #
1806/// #     fn spawn_blocking<F>(f: F) -> Self::JoinHandle where F: FnOnce() + Send + 'static {
1807/// #         unreachable!()
1808/// #     }
1809/// # }
1810/// #
1811/// # impl ContextExt for MyCustomRuntime {
1812/// #     fn scope<F, R>(locals: TaskLocals, fut: F) -> Pin<Box<dyn Future<Output = R> + Send>>
1813/// #     where
1814/// #         F: Future<Output = R> + Send + 'static
1815/// #     {
1816/// #         unreachable!()
1817/// #     }
1818/// #     fn get_task_locals() -> Option<TaskLocals> {
1819/// #         unreachable!()
1820/// #     }
1821/// # }
1822///
1823/// use pyo3::prelude::*;
1824/// use futures_util::stream::{StreamExt, TryStreamExt};
1825/// use std::ffi::CString;
1826///
1827/// const TEST_MOD: &str = r#"
1828/// import asyncio
1829///
1830/// async def gen():
1831///     for i in range(10):
1832///         await asyncio.sleep(0.1)
1833///         yield i
1834/// "#;
1835///
1836/// # async fn test_async_gen() -> PyResult<()> {
1837/// let stream = Python::attach(|py| {
1838///     let test_mod = PyModule::from_code(
1839///         py,
1840///         &CString::new(TEST_MOD).unwrap(),
1841///         &CString::new("test_rust_coroutine/test_mod.py").unwrap(),
1842///         &CString::new("test_mod").unwrap(),
1843///     )?;
1844///
1845///     pyo3_async_runtimes::generic::into_stream_v2::<MyCustomRuntime>(test_mod.call_method0("gen")?)
1846/// })?;
1847///
1848/// let vals = stream
1849///     .map(|item| Python::attach(|py| -> PyResult<i32> { Ok(item.bind(py).extract()?) }))
1850///     .try_collect::<Vec<i32>>()
1851///     .await?;
1852///
1853/// assert_eq!((0..10).collect::<Vec<i32>>(), vals);
1854///
1855/// Ok(())
1856/// # }
1857/// ```
1858#[cfg(feature = "unstable-streams")]
1859pub fn into_stream_v2<R>(
1860    gen: Bound<'_, PyAny>,
1861) -> PyResult<impl futures_util::stream::Stream<Item = Py<PyAny>> + 'static>
1862where
1863    R: Runtime + ContextExt,
1864{
1865    into_stream_with_locals_v2::<R>(get_current_locals::<R>(gen.py())?, gen)
1866}
pyo3_async_runtimes/generic.rs

pyo3_async_runtimes/
generic.rs
