pyo3/coroutine/
cancel.rs

1use crate::{Py, PyAny};
2use alloc::sync::Arc;
3use core::future::poll_fn;
4use core::task::{Context, Poll, Waker};
5use std::sync::Mutex;
6
7#[derive(Debug, Default)]
8struct Inner {
9    exception: Option<Py<PyAny>>,
10    waker: Option<Waker>,
11}
12
13/// Helper used to wait and retrieve exception thrown in [`Coroutine`](super::Coroutine).
14///
15/// Only the last exception thrown can be retrieved.
16#[derive(Debug, Default)]
17pub struct CancelHandle(Arc<Mutex<Inner>>);
18
19impl CancelHandle {
20    /// Create a new `CoroutineCancel`.
21    pub fn new() -> Self {
22        Default::default()
23    }
24
25    /// Returns whether the associated coroutine has been cancelled.
26    pub fn is_cancelled(&self) -> bool {
27        self.0.lock().unwrap().exception.is_some()
28    }
29
30    /// Poll to retrieve the exception thrown in the associated coroutine.
31    pub fn poll_cancelled(&mut self, cx: &mut Context<'_>) -> Poll<Py<PyAny>> {
32        let mut inner = self.0.lock().unwrap();
33        if let Some(exc) = inner.exception.take() {
34            return Poll::Ready(exc);
35        }
36        if let Some(ref waker) = inner.waker {
37            if cx.waker().will_wake(waker) {
38                return Poll::Pending;
39            }
40        }
41        inner.waker = Some(cx.waker().clone());
42        Poll::Pending
43    }
44
45    /// Retrieve the exception thrown in the associated coroutine.
46    pub async fn cancelled(&mut self) -> Py<PyAny> {
47        poll_fn(|cx| self.poll_cancelled(cx)).await
48    }
49
50    #[doc(hidden)]
51    pub fn throw_callback(&self) -> ThrowCallback {
52        ThrowCallback(self.0.clone())
53    }
54}
55
56#[doc(hidden)]
57pub struct ThrowCallback(Arc<Mutex<Inner>>);
58
59impl ThrowCallback {
60    pub(super) fn throw(&self, exc: Py<PyAny>) {
61        let mut inner = self.0.lock().unwrap();
62        inner.exception = Some(exc);
63        if let Some(waker) = inner.waker.take() {
64            waker.wake();
65        }
66    }
67}
pyo3/coroutine/cancel.rs

pyo3/coroutine/
cancel.rs
