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, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine,
8    class::PyClassImpl,
9    coroutine::{Coro, warn_deprecated_throw_signature},
10    frame::FrameObjectRef,
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    // Tracked in `make_generator_or_coro`, together with the frame the object
31    // is born owning, so the pair costs one trip through the GC's gen0 list
32    // instead of two.
33    const NEW_REF_UNTRACKED: bool = true;
34
35    #[inline]
36    fn class(ctx: &Context) -> &'static Py<PyType> {
37        ctx.types.generator_type
38    }
39}
40
41impl PyGenerator {
42    pub const fn as_coro(&self) -> &Coro {
43        &self.inner
44    }
45
46    #[must_use]
47    pub fn new(frame: FrameObjectRef, name: PyStrRef, qualname: PyStrRef) -> Self {
48        Self {
49            inner: Coro::new(frame, name, qualname),
50        }
51    }
52}
53
54#[pyclass(
55    itemsize = core::mem::size_of::<crate::PyObjectRef>(),
56    flags(DISALLOW_INSTANTIATION, HAS_WEAKREF),
57    with(Py, IterNext, Iterable, Representable, Destructor)
58)]
59impl PyGenerator {}
60
61#[pyclass]
62impl Py<PyGenerator> {
63    #[pymethod]
64    fn send(&self, value: PyObjectRef, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
65        self.inner.send(self.as_object(), value, vm)
66    }
67
68    #[pymethod]
69    fn throw(
70        &self,
71        exc_type: PyObjectRef,
72        exc_val: OptionalArg,
73        exc_tb: OptionalArg,
74        vm: &VirtualMachine,
75    ) -> PyResult<PyIterReturn> {
76        warn_deprecated_throw_signature(&exc_val, &exc_tb, vm)?;
77        self.inner.throw(
78            self.as_object(),
79            exc_type,
80            exc_val.unwrap_or_none(vm),
81            exc_tb.unwrap_or_none(vm),
82            vm,
83        )
84    }
85
86    #[pymethod]
87    fn close(&self, vm: &VirtualMachine) -> PyResult<PyObjectRef> {
88        self.inner.close(self.as_object(), vm)
89    }
90
91    #[pygetset]
92    fn __name__(&self) -> PyStrRef {
93        self.inner.name()
94    }
95
96    #[pygetset(setter)]
97    fn set___name__(&self, name: PyStrRef) {
98        self.inner.set_name(name)
99    }
100
101    #[pygetset]
102    fn __qualname__(&self) -> PyStrRef {
103        self.inner.qualname()
104    }
105
106    #[pygetset(setter)]
107    fn set___qualname__(&self, qualname: PyStrRef) {
108        self.inner.set_qualname(qualname)
109    }
110
111    #[pygetset]
112    fn gi_frame(&self, _vm: &VirtualMachine) -> Option<FrameObjectRef> {
113        if self.inner.closed() {
114            None
115        } else {
116            self.inner.frame_opt()
117        }
118    }
119
120    #[pygetset]
121    fn gi_running(&self, _vm: &VirtualMachine) -> bool {
122        self.inner.running()
123    }
124
125    #[pygetset]
126    fn gi_code(&self, _vm: &VirtualMachine) -> PyRef<PyCode> {
127        self.inner.code()
128    }
129
130    #[pygetset]
131    fn gi_yieldfrom(&self, _vm: &VirtualMachine) -> Option<PyObjectRef> {
132        self.inner.frame_opt().and_then(|f| f.yield_from_target())
133    }
134
135    #[pygetset]
136    fn gi_suspended(&self, _vm: &VirtualMachine) -> bool {
137        self.inner.suspended()
138    }
139
140    #[pyclassmethod]
141    fn __class_getitem__(
142        cls: PyTypeRef,
143        args: PyObjectRef,
144        vm: &VirtualMachine,
145    ) -> PyResult<PyGenericAlias> {
146        PyGenericAlias::from_args(cls, args, vm)
147    }
148}
149
150impl Representable for PyGenerator {
151    #[inline]
152    fn repr_str(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<String> {
153        Ok(zelf.inner.repr(zelf.as_object(), zelf.get_id(), vm))
154    }
155}
156
157impl SelfIter for PyGenerator {}
158impl IterNext for PyGenerator {
159    fn next(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
160        zelf.send(vm.ctx.none(), vm)
161    }
162}
163
164impl Destructor for PyGenerator {
165    fn del(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<()> {
166        // _PyGen_Finalize: close the generator if it's still suspended
167        if zelf.inner.closed() || zelf.inner.running() {
168            return Ok(());
169        }
170        if let Err(e) = zelf.inner.close(zelf.as_object(), vm) {
171            crate::coroutine::unraisable_while_closing(zelf.as_object(), &zelf.inner, e, vm);
172        }
173        Ok(())
174    }
175}
176
177impl Drop for PyGenerator {
178    fn drop(&mut self) {
179        if let Some(frame) = self.inner.frame_opt() {
180            frame.clear_generator();
181        }
182    }
183}
184
185/// Fast, VM-free check mirroring the read-only-state branches of
186/// `<PyGenerator as Destructor>::del`: a generator that's already closed or
187/// currently running needs no `close()`-style cleanup, so its `del` slot is a
188/// documented no-op. Skipping the call avoids attaching to a VM
189/// (`with_vm`) on every generator drop for the common case of a generator
190/// consumed to completion.
191fn generator_del_needed(zelf: &PyObject) -> bool {
192    let zelf: &Py<PyGenerator> = zelf
193        .downcast_ref()
194        .expect("del_needed is only installed on the generator type");
195    !(zelf.inner.closed() || zelf.inner.running())
196}
197
198pub(crate) fn init(ctx: &'static Context) {
199    PyGenerator::extend_class(ctx, ctx.types.generator_type);
200    ctx.types
201        .generator_type
202        .slots
203        .del_needed
204        .store(Some(generator_del_needed));
205}
rustpython_vm/builtins/generator.rs

rustpython_vm/builtins/
generator.rs
