pyo3_pytests/
awaitable.rs

1//! The following classes are examples of objects which implement Python's
2//! awaitable protocol.
3//!
4//! Both IterAwaitable and FutureAwaitable will return a value immediately
5//! when awaited, see guide examples related to pyo3-async-runtimes for ways
6//! to suspend tasks and await results.
7
8use pyo3::exceptions::PyStopIteration;
9use pyo3::prelude::*;
10
11#[pyclass]
12#[derive(Debug)]
13pub(crate) struct IterAwaitable {
14    result: Option<PyResult<PyObject>>,
15}
16
17#[pymethods]
18impl IterAwaitable {
19    #[new]
20    fn new(result: PyObject) -> Self {
21        IterAwaitable {
22            result: Some(Ok(result)),
23        }
24    }
25
26    fn __await__(pyself: PyRef<'_, Self>) -> PyRef<'_, Self> {
27        pyself
28    }
29
30    fn __iter__(pyself: PyRef<'_, Self>) -> PyRef<'_, Self> {
31        pyself
32    }
33
34    fn __next__(&mut self, py: Python<'_>) -> PyResult<PyObject> {
35        match self.result.take() {
36            Some(res) => match res {
37                Ok(v) => Err(PyStopIteration::new_err(v)),
38                Err(err) => Err(err),
39            },
40            _ => Ok(py.None()),
41        }
42    }
43}
44
45#[pyclass]
46pub(crate) struct FutureAwaitable {
47    #[pyo3(get, set, name = "_asyncio_future_blocking")]
48    py_block: bool,
49    result: Option<PyResult<PyObject>>,
50}
51
52#[pymethods]
53impl FutureAwaitable {
54    #[new]
55    fn new(result: PyObject) -> Self {
56        FutureAwaitable {
57            py_block: false,
58            result: Some(Ok(result)),
59        }
60    }
61
62    fn __await__(pyself: PyRef<'_, Self>) -> PyRef<'_, Self> {
63        pyself
64    }
65
66    fn __iter__(pyself: PyRef<'_, Self>) -> PyRef<'_, Self> {
67        pyself
68    }
69
70    fn __next__(mut pyself: PyRefMut<'_, Self>) -> PyResult<PyRefMut<'_, Self>> {
71        match pyself.result {
72            Some(_) => match pyself.result.take().unwrap() {
73                Ok(v) => Err(PyStopIteration::new_err(v)),
74                Err(err) => Err(err),
75            },
76            _ => Ok(pyself),
77        }
78    }
79}
80
81#[pymodule(gil_used = false)]
82pub fn awaitable(m: &Bound<'_, PyModule>) -> PyResult<()> {
83    m.add_class::<IterAwaitable>()?;
84    m.add_class::<FutureAwaitable>()?;
85    Ok(())
86}
pyo3_pytests/awaitable.rs

pyo3_pytests/
awaitable.rs
