rustpython_vm/builtins/
coroutine.rs

1use super::{PyCode, PyStrRef, PyType};
2use crate::{
3    class::PyClassImpl,
4    coroutine::Coro,
5    frame::FrameRef,
6    function::OptionalArg,
7    protocol::PyIterReturn,
8    types::{IterNext, Iterable, Representable, SelfIter, Unconstructible},
9    AsObject, Context, Py, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine,
10};
11
12#[pyclass(module = false, name = "coroutine")]
13#[derive(Debug)]
14// PyCoro_Type in CPython
15pub struct PyCoroutine {
16    inner: Coro,
17}
18
19impl PyPayload for PyCoroutine {
20    fn class(ctx: &Context) -> &'static Py<PyType> {
21        ctx.types.coroutine_type
22    }
23}
24
25#[pyclass(with(Py, Unconstructible, IterNext, Representable))]
26impl PyCoroutine {
27    pub fn as_coro(&self) -> &Coro {
28        &self.inner
29    }
30
31    pub fn new(frame: FrameRef, name: PyStrRef) -> Self {
32        PyCoroutine {
33            inner: Coro::new(frame, name),
34        }
35    }
36
37    #[pygetset(magic)]
38    fn name(&self) -> PyStrRef {
39        self.inner.name()
40    }
41
42    #[pygetset(magic, setter)]
43    fn set_name(&self, name: PyStrRef) {
44        self.inner.set_name(name)
45    }
46
47    #[pymethod(name = "__await__")]
48    fn r#await(zelf: PyRef<Self>) -> PyCoroutineWrapper {
49        PyCoroutineWrapper { coro: zelf }
50    }
51
52    #[pygetset]
53    fn cr_await(&self, _vm: &VirtualMachine) -> Option<PyObjectRef> {
54        self.inner.frame().yield_from_target()
55    }
56    #[pygetset]
57    fn cr_frame(&self, _vm: &VirtualMachine) -> FrameRef {
58        self.inner.frame()
59    }
60    #[pygetset]
61    fn cr_running(&self, _vm: &VirtualMachine) -> bool {
62        self.inner.running()
63    }
64    #[pygetset]
65    fn cr_code(&self, _vm: &VirtualMachine) -> PyRef<PyCode> {
66        self.inner.frame().code.clone()
67    }
68    // TODO: coroutine origin tracking:
69    // https://docs.python.org/3/library/sys.html#sys.set_coroutine_origin_tracking_depth
70    #[pygetset]
71    fn cr_origin(&self, _vm: &VirtualMachine) -> Option<(PyStrRef, usize, PyStrRef)> {
72        None
73    }
74}
75
76#[pyclass]
77impl Py<PyCoroutine> {
78    #[pymethod]
79    fn send(&self, value: PyObjectRef, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
80        self.inner.send(self.as_object(), value, vm)
81    }
82
83    #[pymethod]
84    fn throw(
85        &self,
86        exc_type: PyObjectRef,
87        exc_val: OptionalArg,
88        exc_tb: OptionalArg,
89        vm: &VirtualMachine,
90    ) -> PyResult<PyIterReturn> {
91        self.inner.throw(
92            self.as_object(),
93            exc_type,
94            exc_val.unwrap_or_none(vm),
95            exc_tb.unwrap_or_none(vm),
96            vm,
97        )
98    }
99
100    #[pymethod]
101    fn close(&self, vm: &VirtualMachine) -> PyResult<()> {
102        self.inner.close(self.as_object(), vm)
103    }
104}
105
106impl Unconstructible for PyCoroutine {}
107
108impl Representable for PyCoroutine {
109    #[inline]
110    fn repr_str(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<String> {
111        Ok(zelf.inner.repr(zelf.as_object(), zelf.get_id(), vm))
112    }
113}
114
115impl SelfIter for PyCoroutine {}
116impl IterNext for PyCoroutine {
117    fn next(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
118        zelf.send(vm.ctx.none(), vm)
119    }
120}
121
122#[pyclass(module = false, name = "coroutine_wrapper")]
123#[derive(Debug)]
124// PyCoroWrapper_Type in CPython
125pub struct PyCoroutineWrapper {
126    coro: PyRef<PyCoroutine>,
127}
128
129impl PyPayload for PyCoroutineWrapper {
130    fn class(ctx: &Context) -> &'static Py<PyType> {
131        ctx.types.coroutine_wrapper_type
132    }
133}
134
135#[pyclass(with(IterNext, Iterable))]
136impl PyCoroutineWrapper {
137    #[pymethod]
138    fn send(&self, val: PyObjectRef, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
139        self.coro.send(val, vm)
140    }
141
142    #[pymethod]
143    fn throw(
144        &self,
145        exc_type: PyObjectRef,
146        exc_val: OptionalArg,
147        exc_tb: OptionalArg,
148        vm: &VirtualMachine,
149    ) -> PyResult<PyIterReturn> {
150        self.coro.throw(exc_type, exc_val, exc_tb, vm)
151    }
152}
153
154impl SelfIter for PyCoroutineWrapper {}
155impl IterNext for PyCoroutineWrapper {
156    fn next(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
157        Self::send(zelf, vm.ctx.none(), vm)
158    }
159}
160
161pub fn init(ctx: &Context) {
162    PyCoroutine::extend_class(ctx, ctx.types.coroutine_type);
163    PyCoroutineWrapper::extend_class(ctx, ctx.types.coroutine_wrapper_type);
164}
rustpython_vm/builtins/coroutine.rs

rustpython_vm/builtins/
coroutine.rs
