rustpython_vm/builtins/
generator.rs

1/*
2 * The mythical generator.
3 */
4
5use super::{PyCode, PyGenericAlias, PyStrRef, PyType, PyTypeRef};
6use crate::{
7    AsObject, Context, Py, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine,
8    class::PyClassImpl,
9    coroutine::{Coro, warn_deprecated_throw_signature},
10    frame::FrameRef,
11    function::OptionalArg,
12    object::{Traverse, TraverseFn},
13    protocol::PyIterReturn,
14    types::{Destructor, IterNext, Iterable, Representable, SelfIter},
15};
16
17#[pyclass(module = false, name = "generator", traverse = "manual")]
18#[derive(Debug)]
19pub struct PyGenerator {
20    inner: Coro,
21}
22
23unsafe impl Traverse for PyGenerator {
24    fn traverse(&self, tracer_fn: &mut TraverseFn<'_>) {
25        self.inner.traverse(tracer_fn);
26    }
27}
28
29impl PyPayload for PyGenerator {
30    #[inline]
31    fn class(ctx: &Context) -> &'static Py<PyType> {
32        ctx.types.generator_type
33    }
34}
35
36#[pyclass(
37    flags(DISALLOW_INSTANTIATION, HAS_WEAKREF),
38    with(Py, IterNext, Iterable, Representable, Destructor)
39)]
40impl PyGenerator {
41    pub const fn as_coro(&self) -> &Coro {
42        &self.inner
43    }
44
45    pub fn new(frame: FrameRef, name: PyStrRef, qualname: PyStrRef) -> Self {
46        Self {
47            inner: Coro::new(frame, name, qualname),
48        }
49    }
50
51    #[pygetset]
52    fn __name__(&self) -> PyStrRef {
53        self.inner.name()
54    }
55
56    #[pygetset(setter)]
57    fn set___name__(&self, name: PyStrRef) {
58        self.inner.set_name(name)
59    }
60
61    #[pygetset]
62    fn __qualname__(&self) -> PyStrRef {
63        self.inner.qualname()
64    }
65
66    #[pygetset(setter)]
67    fn set___qualname__(&self, qualname: PyStrRef) {
68        self.inner.set_qualname(qualname)
69    }
70
71    #[pygetset]
72    fn gi_frame(&self, _vm: &VirtualMachine) -> Option<FrameRef> {
73        if self.inner.closed() {
74            None
75        } else {
76            Some(self.inner.frame())
77        }
78    }
79
80    #[pygetset]
81    fn gi_running(&self, _vm: &VirtualMachine) -> bool {
82        self.inner.running()
83    }
84
85    #[pygetset]
86    fn gi_code(&self, _vm: &VirtualMachine) -> PyRef<PyCode> {
87        self.inner.frame().code.clone()
88    }
89
90    #[pygetset]
91    fn gi_yieldfrom(&self, _vm: &VirtualMachine) -> Option<PyObjectRef> {
92        self.inner.frame().yield_from_target()
93    }
94
95    #[pygetset]
96    fn gi_suspended(&self, _vm: &VirtualMachine) -> bool {
97        self.inner.suspended()
98    }
99
100    #[pyclassmethod]
101    fn __class_getitem__(cls: PyTypeRef, args: PyObjectRef, vm: &VirtualMachine) -> PyGenericAlias {
102        PyGenericAlias::from_args(cls, args, vm)
103    }
104}
105
106#[pyclass]
107impl Py<PyGenerator> {
108    #[pymethod]
109    fn send(&self, value: PyObjectRef, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
110        self.inner.send(self.as_object(), value, vm)
111    }
112
113    #[pymethod]
114    fn throw(
115        &self,
116        exc_type: PyObjectRef,
117        exc_val: OptionalArg,
118        exc_tb: OptionalArg,
119        vm: &VirtualMachine,
120    ) -> PyResult<PyIterReturn> {
121        warn_deprecated_throw_signature(&exc_val, &exc_tb, vm)?;
122        self.inner.throw(
123            self.as_object(),
124            exc_type,
125            exc_val.unwrap_or_none(vm),
126            exc_tb.unwrap_or_none(vm),
127            vm,
128        )
129    }
130
131    #[pymethod]
132    fn close(&self, vm: &VirtualMachine) -> PyResult<PyObjectRef> {
133        self.inner.close(self.as_object(), vm)
134    }
135}
136
137impl Representable for PyGenerator {
138    #[inline]
139    fn repr_str(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<String> {
140        Ok(zelf.inner.repr(zelf.as_object(), zelf.get_id(), vm))
141    }
142}
143
144impl SelfIter for PyGenerator {}
145impl IterNext for PyGenerator {
146    fn next(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
147        zelf.send(vm.ctx.none(), vm)
148    }
149}
150
151impl Destructor for PyGenerator {
152    fn del(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<()> {
153        // _PyGen_Finalize: close the generator if it's still suspended
154        if zelf.inner.closed() || zelf.inner.running() {
155            return Ok(());
156        }
157        // Generator was never started, just mark as closed
158        if zelf.inner.frame().lasti() == 0 {
159            zelf.inner.closed.store(true);
160            return Ok(());
161        }
162        // Throw GeneratorExit to run finally blocks
163        if let Err(e) = zelf.inner.close(zelf.as_object(), vm) {
164            vm.run_unraisable(e, None, zelf.as_object().to_owned());
165        }
166        Ok(())
167    }
168}
169
170impl Drop for PyGenerator {
171    fn drop(&mut self) {
172        self.inner.frame().clear_generator();
173    }
174}
175
176pub fn init(ctx: &'static Context) {
177    PyGenerator::extend_class(ctx, ctx.types.generator_type);
178}
rustpython_vm/builtins/generator.rs

rustpython_vm/builtins/
generator.rs
