rustpython_vm/builtins/
frame_locals_proxy.rs

1//! The `FrameLocalsProxy` type returned by `frame.f_locals` for optimized
2//! (function) frames and for unoptimized frames that currently have PEP 709
3//! hidden comprehension locals. Implements PEP 667 write-through semantics on
4//! top of the frame's fast-local slots and an extra-locals side dict.
5
6use super::{PyDict, PyDictRef, PyType};
7use crate::{
8    AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine,
9    atomic_func,
10    class::{PyClassDef, PyClassImpl},
11    frame::FrameObjectRef,
12    function::{FuncArgs, OptionalArg, PyArithmeticValue, PyComparisonValue},
13    object::{Traverse, TraverseFn},
14    protocol::{PyIterReturn, PyMappingMethods, PyNumberMethods, PySequenceMethods},
15    recursion::ReprGuard,
16    types::{
17        AsMapping, AsNumber, AsSequence, Comparable, Constructor, Iterable, PyComparisonOp,
18        Representable,
19    },
20};
21use rustpython_common::lock::LazyLock;
22use rustpython_common::wtf8::Wtf8Buf;
23
24#[pyclass(
25    module = false,
26    name = "FrameLocalsProxy",
27    unhashable = true,
28    traverse = "manual"
29)]
30#[derive(Debug)]
31pub struct FrameLocalsProxy {
32    frame: FrameObjectRef,
33}
34
35unsafe impl Traverse for FrameLocalsProxy {
36    fn traverse(&self, tracer_fn: &mut TraverseFn<'_>) {
37        self.frame.traverse(tracer_fn);
38    }
39}
40
41impl PyPayload for FrameLocalsProxy {
42    #[inline]
43    fn class(ctx: &Context) -> &'static Py<PyType> {
44        ctx.types.frame_locals_proxy_type
45    }
46}
47
48impl FrameLocalsProxy {
49    pub(crate) fn new(frame: FrameObjectRef) -> Self {
50        Self { frame }
51    }
52
53    fn snapshot(&self, vm: &VirtualMachine) -> PyResult<PyDictRef> {
54        self.frame.framelocalsproxy_snapshot(vm)
55    }
56
57    fn items_vec(&self, vm: &VirtualMachine) -> PyResult<Vec<(PyObjectRef, PyObjectRef)>> {
58        self.frame.framelocalsproxy_items(vm)
59    }
60
61    fn keys_vec(&self, vm: &VirtualMachine) -> PyResult<Vec<PyObjectRef>> {
62        Ok(self
63            .items_vec(vm)?
64            .into_iter()
65            .map(|(key, _)| key)
66            .collect())
67    }
68}
69
70impl Constructor for FrameLocalsProxy {
71    type Args = FuncArgs;
72
73    fn py_new(_cls: &Py<PyType>, args: Self::Args, vm: &VirtualMachine) -> PyResult<Self> {
74        if args.args.len() != 1 {
75            return Err(vm.new_arity_type_error(Self::NAME, 1..=1, args.args.len()));
76        }
77        if !args.kwargs.is_empty() {
78            return Err(vm.new_type_error("FrameLocalsProxy() takes no keyword arguments"));
79        }
80        let mut args = args.args;
81        let frame: FrameObjectRef = args
82            .pop()
83            .unwrap()
84            .downcast()
85            .map_err(|_| vm.new_type_error("FrameLocalsProxy expected a frame"))?;
86        Ok(Self::new(frame))
87    }
88}
89
90impl FrameLocalsProxy {
91    fn __getitem__(&self, key: PyObjectRef, vm: &VirtualMachine) -> PyResult {
92        self.frame.framelocalsproxy_getitem(key, vm)
93    }
94
95    fn __setitem__(
96        &self,
97        key: PyObjectRef,
98        value: PyObjectRef,
99        vm: &VirtualMachine,
100    ) -> PyResult<()> {
101        self.frame.framelocalsproxy_setitem(&key, value, vm)
102    }
103
104    fn __delitem__(&self, key: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
105        self.frame.framelocalsproxy_delitem(key, vm)
106    }
107
108    fn __contains__(&self, key: PyObjectRef, vm: &VirtualMachine) -> PyResult<bool> {
109        self.frame.framelocalsproxy_contains(&key, vm)
110    }
111
112    fn __len__(&self, vm: &VirtualMachine) -> PyResult<usize> {
113        Ok(self.items_vec(vm)?.len())
114    }
115
116    fn update_from(&self, other: &PyObject, vm: &VirtualMachine) -> PyResult<()> {
117        if other.downcast_ref::<PyDict>().is_none() && other.downcast_ref::<Self>().is_none() {
118            return Err(
119                vm.new_type_error("update() argument must be dict or another FrameLocalsProxy")
120            );
121        }
122        // CPython deliberately uses the mapping protocol here, including
123        // overridden keys()/__getitem__ on dict subclasses.
124        let keys = other
125            .get_attr(vm.ctx.intern_str("keys"), vm)?
126            .call((), vm)?
127            .get_iter(vm)?;
128        while let PyIterReturn::Return(key) = keys.next(vm)? {
129            let value = other.get_item(&*key, vm)?;
130            self.frame.framelocalsproxy_setitem(&key, value, vm)?;
131        }
132        Ok(())
133    }
134
135    fn __ior__(zelf: PyRef<Self>, other: &PyObject, vm: &VirtualMachine) -> PyResult {
136        if other.downcast_ref::<PyDict>().is_none() && other.downcast_ref::<Self>().is_none() {
137            return Ok(vm.ctx.not_implemented());
138        }
139        zelf.update_from(other, vm)?;
140        Ok(zelf.into())
141    }
142
143    fn __or__(&self, other: PyObjectRef, vm: &VirtualMachine) -> PyResult {
144        if other.downcast_ref::<PyDict>().is_none() && other.downcast_ref::<Self>().is_none() {
145            return Ok(vm.ctx.not_implemented());
146        }
147        let result = self.snapshot(vm)?;
148        if other.downcast_ref::<PyDict>().is_some() {
149            // PyDict_Update reads a dict subclass's stored entries directly;
150            // it does not dispatch to overridden mapping methods.
151            result.merge_dict(&other.downcast().unwrap(), true, vm)?;
152        } else {
153            result.merge_object(other, vm)?;
154        }
155        Ok(result.into())
156    }
157
158    fn __ror__(&self, other: &PyObject, vm: &VirtualMachine) -> PyResult {
159        let Some(other) = other.downcast_ref::<PyDict>() else {
160            return Ok(vm.ctx.not_implemented());
161        };
162        let result = other.copy().into_ref(&vm.ctx);
163        result.merge_object(self.snapshot(vm)?.into(), vm)?;
164        Ok(result.into())
165    }
166}
167
168#[pyclass(with(
169    Constructor,
170    AsMapping,
171    AsSequence,
172    AsNumber,
173    Iterable,
174    Comparable,
175    Representable
176))]
177impl Py<FrameLocalsProxy> {
178    #[pymethod]
179    fn keys(&self, vm: &VirtualMachine) -> PyResult<PyObjectRef> {
180        Ok(vm.ctx.new_list(self.keys_vec(vm)?).into())
181    }
182
183    #[pymethod]
184    fn values(&self, vm: &VirtualMachine) -> PyResult<PyObjectRef> {
185        let values = self
186            .items_vec(vm)?
187            .into_iter()
188            .map(|(_, value)| value)
189            .collect();
190        Ok(vm.ctx.new_list(values).into())
191    }
192
193    #[pymethod]
194    fn items(&self, vm: &VirtualMachine) -> PyResult<PyObjectRef> {
195        let items = self
196            .items_vec(vm)?
197            .into_iter()
198            .map(|(k, v)| vm.ctx.new_tuple(vec![k, v]).into())
199            .collect();
200        Ok(vm.ctx.new_list(items).into())
201    }
202
203    #[pymethod]
204    fn get(&self, key: PyObjectRef, default: OptionalArg, vm: &VirtualMachine) -> PyResult {
205        match self.frame.framelocalsproxy_getitem(key, vm) {
206            Ok(value) => Ok(value),
207            Err(e) if e.fast_isinstance(vm.ctx.exceptions.key_error) => {
208                Ok(default.unwrap_or_none(vm))
209            }
210            Err(e) => Err(e),
211        }
212    }
213
214    #[pymethod]
215    fn pop(&self, key: PyObjectRef, default: OptionalArg, vm: &VirtualMachine) -> PyResult {
216        self.frame
217            .framelocalsproxy_pop(key, default.into_option(), vm)
218    }
219
220    #[pymethod]
221    fn setdefault(&self, key: PyObjectRef, default: OptionalArg, vm: &VirtualMachine) -> PyResult {
222        self.frame
223            .framelocalsproxy_setdefault(&key, default.unwrap_or_none(vm), vm)
224    }
225
226    #[pymethod]
227    fn copy(&self, vm: &VirtualMachine) -> PyResult<PyObjectRef> {
228        Ok(self.snapshot(vm)?.into())
229    }
230
231    #[pymethod]
232    fn update(&self, args: FuncArgs, vm: &VirtualMachine) -> PyResult<()> {
233        if !args.kwargs.is_empty() {
234            return Err(vm.new_type_error("FrameLocalsProxy.update() takes no keyword arguments"));
235        }
236        if args.args.len() != 1 {
237            return Err(vm.new_type_error(format!(
238                "FrameLocalsProxy.update() takes exactly one argument ({} given)",
239                args.args.len()
240            )));
241        }
242        self.update_from(&args.args[0], vm)
243    }
244
245    #[pymethod]
246    fn __reversed__(&self, vm: &VirtualMachine) -> PyResult<PyObjectRef> {
247        let mut keys = self.keys_vec(vm)?;
248        keys.reverse();
249        Ok(vm.ctx.new_list(keys).into())
250    }
251
252    #[pymethod]
253    fn __reduce__(&self, vm: &VirtualMachine) -> PyResult {
254        Err(vm.new_type_error("cannot pickle 'FrameLocalsProxy' object"))
255    }
256
257    #[pymethod]
258    fn __reduce_ex__(&self, _protocol: OptionalArg, vm: &VirtualMachine) -> PyResult {
259        Err(vm.new_type_error("cannot pickle 'FrameLocalsProxy' object"))
260    }
261}
262
263impl AsMapping for FrameLocalsProxy {
264    fn as_mapping() -> &'static PyMappingMethods {
265        static AS_MAPPING: LazyLock<PyMappingMethods> = LazyLock::new(|| PyMappingMethods {
266            length: atomic_func!(
267                |mapping, vm| FrameLocalsProxy::mapping_downcast(mapping).__len__(vm)
268            ),
269            subscript: atomic_func!(|mapping, needle, vm| {
270                FrameLocalsProxy::mapping_downcast(mapping).__getitem__(needle.to_owned(), vm)
271            }),
272            ass_subscript: atomic_func!(|mapping, needle, value, vm| {
273                let zelf = FrameLocalsProxy::mapping_downcast(mapping);
274                match value {
275                    Some(value) => zelf.__setitem__(needle.to_owned(), value, vm),
276                    None => zelf.__delitem__(needle.to_owned(), vm),
277                }
278            }),
279        });
280        &AS_MAPPING
281    }
282}
283
284impl AsSequence for FrameLocalsProxy {
285    fn as_sequence() -> &'static PySequenceMethods {
286        static AS_SEQUENCE: LazyLock<PySequenceMethods> = LazyLock::new(|| PySequenceMethods {
287            contains: atomic_func!(|seq, target, vm| {
288                FrameLocalsProxy::sequence_downcast(seq).__contains__(target.to_owned(), vm)
289            }),
290            ..PySequenceMethods::NOT_IMPLEMENTED
291        });
292        &AS_SEQUENCE
293    }
294}
295
296impl AsNumber for FrameLocalsProxy {
297    fn as_number() -> &'static PyNumberMethods {
298        static AS_NUMBER: PyNumberMethods = PyNumberMethods {
299            or: Some(|a, b, vm| {
300                if let Some(proxy) = a.downcast_ref::<FrameLocalsProxy>() {
301                    proxy.__or__(b.to_owned(), vm)
302                } else if let Some(proxy) = b.downcast_ref::<FrameLocalsProxy>() {
303                    proxy.__ror__(a, vm)
304                } else {
305                    Ok(vm.ctx.not_implemented())
306                }
307            }),
308            inplace_or: Some(|a, b, vm| {
309                let proxy = a
310                    .to_owned()
311                    .downcast::<FrameLocalsProxy>()
312                    .map_err(|_| vm.new_type_error("expected FrameLocalsProxy"))?;
313                FrameLocalsProxy::__ior__(proxy, b, vm)
314            }),
315            ..PyNumberMethods::NOT_IMPLEMENTED
316        };
317        &AS_NUMBER
318    }
319}
320
321impl Iterable for FrameLocalsProxy {
322    fn iter(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyResult {
323        let keys = vm.ctx.new_list(zelf.keys_vec(vm)?);
324        keys.as_object().to_owned().get_iter(vm).map(Into::into)
325    }
326}
327
328impl Comparable for FrameLocalsProxy {
329    fn cmp(
330        zelf: &Py<Self>,
331        other: &PyObject,
332        op: PyComparisonOp,
333        vm: &VirtualMachine,
334    ) -> PyResult<PyComparisonValue> {
335        op.eq_only(|| {
336            if let Some(other) = other.downcast_ref::<Self>() {
337                return Ok(PyComparisonValue::Implemented(zelf.frame.is(&other.frame)));
338            }
339            if other.downcast_ref::<PyDict>().is_none() {
340                return Ok(PyComparisonValue::NotImplemented);
341            }
342            let self_dict: PyObjectRef = zelf.snapshot(vm)?.into();
343            let res = self_dict.rich_compare(other.to_owned(), PyComparisonOp::Eq, vm)?;
344            PyArithmeticValue::from_object(vm, res)
345                .map(|o| o.try_to_bool(vm))
346                .transpose()
347        })
348    }
349}
350
351impl Representable for FrameLocalsProxy {
352    fn repr_wtf8(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<Wtf8Buf> {
353        if let Some(_guard) = ReprGuard::enter(vm, zelf.as_object()) {
354            let dict = zelf.snapshot(vm)?;
355            Ok(dict.as_object().repr(vm)?.as_wtf8().to_owned())
356        } else {
357            Ok(Wtf8Buf::from("{...}"))
358        }
359    }
360}
361
362pub(crate) fn init(context: &'static Context) {
363    FrameLocalsProxy::extend_class(context, context.types.frame_locals_proxy_type);
364}
rustpython_vm/builtins/frame_locals_proxy.rs

rustpython_vm/builtins/
frame_locals_proxy.rs
