1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86

//! The following classes are examples of objects which implement Python's
//! awaitable protocol.
//!
//! Both IterAwaitable and FutureAwaitable will return a value immediately
//! when awaited, see guide examples related to pyo3-asyncio for ways
//! to suspend tasks and await results.

use pyo3::exceptions::PyStopIteration;
use pyo3::prelude::*;

#[pyclass]
#[derive(Debug)]
pub(crate) struct IterAwaitable {
    result: Option<PyResult<PyObject>>,
}

#[pymethods]
impl IterAwaitable {
    #[new]
    fn new(result: PyObject) -> Self {
        IterAwaitable {
            result: Some(Ok(result)),
        }
    }

    fn __await__(pyself: PyRef<'_, Self>) -> PyRef<'_, Self> {
        pyself
    }

    fn __iter__(pyself: PyRef<'_, Self>) -> PyRef<'_, Self> {
        pyself
    }

    fn __next__(&mut self, py: Python<'_>) -> PyResult<PyObject> {
        match self.result.take() {
            Some(res) => match res {
                Ok(v) => Err(PyStopIteration::new_err(v)),
                Err(err) => Err(err),
            },
            _ => Ok(py.None()),
        }
    }
}

#[pyclass]
pub(crate) struct FutureAwaitable {
    #[pyo3(get, set, name = "_asyncio_future_blocking")]
    py_block: bool,
    result: Option<PyResult<PyObject>>,
}

#[pymethods]
impl FutureAwaitable {
    #[new]
    fn new(result: PyObject) -> Self {
        FutureAwaitable {
            py_block: false,
            result: Some(Ok(result)),
        }
    }

    fn __await__(pyself: PyRef<'_, Self>) -> PyRef<'_, Self> {
        pyself
    }

    fn __iter__(pyself: PyRef<'_, Self>) -> PyRef<'_, Self> {
        pyself
    }

    fn __next__(mut pyself: PyRefMut<'_, Self>) -> PyResult<PyRefMut<'_, Self>> {
        match pyself.result {
            Some(_) => match pyself.result.take().unwrap() {
                Ok(v) => Err(PyStopIteration::new_err(v)),
                Err(err) => Err(err),
            },
            _ => Ok(pyself),
        }
    }
}

#[pymodule]
pub fn awaitable(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<IterAwaitable>()?;
    m.add_class::<FutureAwaitable>()?;
    Ok(())
}