Skip to content

Commit 067bb73

Browse files
DimitrisJimDimitrisJim
committed
Use PyMapping, fix handling of __prepare__.
1 parent dce4582 commit 067bb73

File tree

6 files changed

+88
-58
lines changed

6 files changed

+88
-58
lines changed

‎vm/src/builtins/frame.rs

+2-1
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
use super::{PyCode, PyDictRef, PyStrRef};
66
use crate::{
77
frame::{Frame, FrameRef},
8+
protocol::PyMapping,
89
types::{Constructor, Unconstructible},
910
IdProtocol, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, VirtualMachine,
1011
};
@@ -44,7 +45,7 @@ impl FrameRef {
4445
}
4546

4647
#[pyproperty]
47-
fn f_locals(self, vm: &VirtualMachine) -> PyResult<PyObjectRef> {
48+
fn f_locals(self, vm: &VirtualMachine) -> PyResult<PyMapping> {
4849
self.locals(vm)
4950
}
5051

‎vm/src/builtins/function.rs

+7-5
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ use crate::{
1010
common::lock::PyMutex,
1111
frame::Frame,
1212
function::{FuncArgs, OptionalArg},
13+
protocol::PyMapping,
1314
scope::Scope,
1415
types::{Callable, Comparable, Constructor, GetAttr, GetDescriptor, PyComparisonOp},
1516
IdProtocol, ItemProtocol, PyClassImpl, PyComparisonValue, PyContext, PyObjectRef, PyRef,
@@ -267,7 +268,7 @@ impl PyFunction {
267268
pub fn invoke_with_locals(
268269
&self,
269270
func_args: FuncArgs,
270-
locals: Option<PyDictRef>,
271+
locals: Option<PyMapping>,
271272
vm: &VirtualMachine,
272273
) -> PyResult {
273274
#[cfg(feature = "jit")]
@@ -287,11 +288,12 @@ impl PyFunction {
287288
let code = &self.code;
288289

289290
let locals = if self.code.flags.contains(bytecode::CodeFlags::NEW_LOCALS) {
290-
vm.ctx.new_dict()
291+
PyMapping::new(vm.ctx.new_dict().into())
292+
} else if let Some(locals) = locals {
293+
locals
291294
} else {
292-
locals.unwrap_or_else(|| self.globals.clone())
293-
}
294-
.into();
295+
PyMapping::new(self.globals.clone().into())
296+
};
295297

296298
// Construct frame:
297299
let frame = Frame::new(

‎vm/src/frame.rs

+40-19
Original file line numberDiff line numberDiff line change
@@ -11,12 +11,12 @@ use crate::{
1111
coroutine::Coro,
1212
exceptions::ExceptionCtor,
1313
function::{FuncArgs, IntoPyResult},
14-
protocol::{PyIter, PyIterReturn},
14+
protocol::{PyIter, PyIterReturn, PyMapping},
1515
scope::Scope,
1616
stdlib::builtins,
1717
types::PyComparisonOp,
18-
IdProtocol, ItemProtocol, PyMethod, PyObjectRef, PyRef, PyResult, PyValue, TryFromObject,
19-
TypeProtocol, VirtualMachine,
18+
IdProtocol, ItemProtocol, PyMethod, PyObjectRef, PyObjectWrap, PyRef, PyResult, PyValue,
19+
TryFromObject, TypeProtocol, VirtualMachine,
2020
};
2121
use indexmap::IndexMap;
2222
use itertools::Itertools;
@@ -99,7 +99,7 @@ pub struct Frame {
9999

100100
pub fastlocals: PyMutex<Box<[Option<PyObjectRef>]>>,
101101
pub(crate) cells_frees: Box<[PyCellRef]>,
102-
pub locals: PyObjectRef,
102+
pub locals: PyMapping,
103103
pub globals: PyDictRef,
104104
pub builtins: PyDictRef,
105105

@@ -179,7 +179,7 @@ impl FrameRef {
179179
f(exec)
180180
}
181181

182-
pub fn locals(&self, vm: &VirtualMachine) -> PyResult<PyObjectRef> {
182+
pub fn locals(&self, vm: &VirtualMachine) -> PyResult<PyMapping> {
183183
let locals = &self.locals;
184184
let code = &**self.code;
185185
let map = &code.varnames;
@@ -188,9 +188,9 @@ impl FrameRef {
188188
let fastlocals = self.fastlocals.lock();
189189
for (k, v) in itertools::zip(&map[..j], &**fastlocals) {
190190
if let Some(v) = v {
191-
locals.set_item(k.clone(), v.clone(), vm)?;
191+
locals.as_object().set_item(k.clone(), v.clone(), vm)?;
192192
} else {
193-
match locals.del_item(k.clone(), vm) {
193+
match locals.as_object().del_item(k.clone(), vm) {
194194
Ok(()) => {}
195195
Err(e) if e.isinstance(&vm.ctx.exceptions.key_error) => {}
196196
Err(e) => return Err(e),
@@ -202,9 +202,9 @@ impl FrameRef {
202202
let map_to_dict = |keys: &[PyStrRef], values: &[PyCellRef]| {
203203
for (k, v) in itertools::zip(keys, values) {
204204
if let Some(v) = v.get() {
205-
locals.set_item(k.clone(), v, vm)?;
205+
locals.as_object().set_item(k.clone(), v, vm)?;
206206
} else {
207-
match locals.del_item(k.clone(), vm) {
207+
match locals.as_object().del_item(k.clone(), vm) {
208208
Ok(()) => {}
209209
Err(e) if e.isinstance(&vm.ctx.exceptions.key_error) => {}
210210
Err(e) => return Err(e),
@@ -275,7 +275,7 @@ struct ExecutingFrame<'a> {
275275
code: &'a PyRef<PyCode>,
276276
fastlocals: &'a PyMutex<Box<[Option<PyObjectRef>]>>,
277277
cells_frees: &'a [PyCellRef],
278-
locals: &'a PyObjectRef,
278+
locals: &'a PyMapping,
279279
globals: &'a PyDictRef,
280280
builtins: &'a PyDictRef,
281281
object: &'a FrameRef,
@@ -497,7 +497,12 @@ impl ExecutingFrame<'_> {
497497
bytecode::Instruction::LoadNameAny(idx) => {
498498
let name = &self.code.names[*idx as usize];
499499
// Try using locals as dict first, if not, fallback to generic method.
500-
let x = match self.locals.clone().downcast_exact::<PyDict>(vm) {
500+
let x = match self
501+
.locals
502+
.clone()
503+
.into_object()
504+
.downcast_exact::<PyDict>(vm)
505+
{
501506
Ok(d) => d.get_item_option(name.clone(), vm)?,
502507
Err(o) => o.get_item(name.clone(), vm).ok(),
503508
};
@@ -525,7 +530,12 @@ impl ExecutingFrame<'_> {
525530
let i = *i as usize;
526531
let name = self.code.freevars[i - self.code.cellvars.len()].clone();
527532
// Try using locals as dict first, if not, fallback to generic method.
528-
let value = match self.locals.clone().downcast_exact::<PyDict>(vm) {
533+
let value = match self
534+
.locals
535+
.clone()
536+
.into_object()
537+
.downcast_exact::<PyDict>(vm)
538+
{
529539
Ok(d) => d.get_item_option(name, vm)?,
530540
Err(o) => o.get_item(name, vm).ok(),
531541
};
@@ -544,8 +554,11 @@ impl ExecutingFrame<'_> {
544554
}
545555
bytecode::Instruction::StoreLocal(idx) => {
546556
let value = self.pop_value();
547-
self.locals
548-
.set_item(self.code.names[*idx as usize].clone(), value, vm)?;
557+
self.locals.as_object().set_item(
558+
self.code.names[*idx as usize].clone(),
559+
value,
560+
vm,
561+
)?;
549562
Ok(None)
550563
}
551564
bytecode::Instruction::StoreGlobal(idx) => {
@@ -565,7 +578,7 @@ impl ExecutingFrame<'_> {
565578
}
566579
bytecode::Instruction::DeleteLocal(idx) => {
567580
let name = &self.code.names[*idx as usize];
568-
match self.locals.del_item(name.clone(), vm) {
581+
match self.locals.clone().into_object().del_item(name.clone(), vm) {
569582
Ok(()) => {}
570583
Err(e) if e.isinstance(&vm.ctx.exceptions.key_error) => {
571584
return Err(vm.new_name_error(format!("name '{}' is not defined", name)))
@@ -724,16 +737,24 @@ impl ExecutingFrame<'_> {
724737
bytecode::Instruction::YieldFrom => self.execute_yield_from(vm),
725738
bytecode::Instruction::SetupAnnotation => {
726739
// Try using locals as dict first, if not, fallback to generic method.
727-
let has_annotations = match self.locals.clone().downcast_exact::<PyDict>(vm) {
740+
let has_annotations = match self
741+
.locals
742+
.clone()
743+
.into_object()
744+
.downcast_exact::<PyDict>(vm)
745+
{
728746
Ok(d) => d.contains_key("__annotations__", vm),
729747
Err(o) => {
730748
let needle = vm.new_pyobj("__annotations__");
731749
self._in(vm, needle, o)?
732750
}
733751
};
734752
if !has_annotations {
735-
self.locals
736-
.set_item("__annotations__", vm.ctx.new_dict().into(), vm)?;
753+
self.locals.as_object().set_item(
754+
"__annotations__",
755+
vm.ctx.new_dict().into(),
756+
vm,
757+
)?;
737758
}
738759
Ok(None)
739760
}
@@ -1146,7 +1167,7 @@ impl ExecutingFrame<'_> {
11461167
for (k, v) in &dict {
11471168
let k = PyStrRef::try_from_object(vm, k)?;
11481169
if filter_pred(k.as_str()) {
1149-
self.locals.set_item(k, v, vm)?;
1170+
self.locals.as_object().set_item(k, v, vm)?;
11501171
}
11511172
}
11521173
}

‎vm/src/scope.rs

+6-5
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,14 @@
11
use crate::{
22
builtins::{pystr::IntoPyStrRef, PyDictRef, PyStrRef},
33
function::IntoPyObject,
4-
ItemProtocol, PyObjectRef, VirtualMachine,
4+
protocol::PyMapping,
5+
ItemProtocol, VirtualMachine,
56
};
67
use std::fmt;
78

89
#[derive(Clone)]
910
pub struct Scope {
10-
pub locals: PyObjectRef,
11+
pub locals: PyMapping,
1112
pub globals: PyDictRef,
1213
}
1314

@@ -20,13 +21,13 @@ impl fmt::Debug for Scope {
2021

2122
impl Scope {
2223
#[inline]
23-
pub fn new(locals: Option<PyObjectRef>, globals: PyDictRef) -> Scope {
24-
let locals = locals.unwrap_or_else(|| globals.clone().into());
24+
pub fn new(locals: Option<PyMapping>, globals: PyDictRef) -> Scope {
25+
let locals = locals.unwrap_or_else(|| PyMapping::new(globals.clone().into()));
2526
Scope { locals, globals }
2627
}
2728

2829
pub fn with_builtins(
29-
locals: Option<PyObjectRef>,
30+
locals: Option<PyMapping>,
3031
globals: PyDictRef,
3132
vm: &VirtualMachine,
3233
) -> Scope {

‎vm/src/stdlib/builtins.rs

+29-24
Original file line numberDiff line numberDiff line change
@@ -32,8 +32,8 @@ mod builtins {
3232
stdlib::sys,
3333
types::PyComparisonOp,
3434
utils::Either,
35-
IdProtocol, ItemProtocol, PyArithmeticValue, PyClassImpl, PyObjectRef, PyRef, PyResult,
36-
PyValue, TryFromObject, TypeProtocol, VirtualMachine,
35+
IdProtocol, ItemProtocol, PyArithmeticValue, PyClassImpl, PyObjectRef, PyObjectWrap, PyRef,
36+
PyResult, PyValue, TryFromObject, TypeProtocol, VirtualMachine,
3737
};
3838
use num_traits::{Signed, ToPrimitive, Zero};
3939

@@ -204,37 +204,33 @@ mod builtins {
204204
struct ScopeArgs {
205205
#[pyarg(any, default)]
206206
globals: Option<PyDictRef>,
207-
// TODO: support any mapping for `locals`
208207
#[pyarg(any, default)]
209-
locals: Option<PyObjectRef>,
208+
locals: Option<PyMapping>,
210209
}
211210

212211
#[cfg(feature = "rustpython-compiler")]
213212
impl ScopeArgs {
214213
fn make_scope(self, vm: &VirtualMachine) -> PyResult<Scope> {
215-
// TODO: Other places where user supplies locals object?
216-
let locals = self.locals;
217-
let get_locals = |default| {
218-
locals.map_or(Ok(default), |locals| {
219-
if !PyMapping::check(&locals, vm) {
220-
Err(vm.new_type_error("locals must be a mapping".to_owned()))
221-
} else {
222-
Ok(locals)
223-
}
224-
})
225-
};
226214
let (globals, locals) = match self.globals {
227215
Some(globals) => {
228216
if !globals.contains_key("__builtins__", vm) {
229217
let builtins_dict = vm.builtins.dict().unwrap().into();
230218
globals.set_item("__builtins__", builtins_dict, vm)?;
231219
}
232-
(globals.clone(), get_locals(globals.into())?)
233-
}
234-
None => {
235-
let globals = vm.current_globals().clone();
236-
(globals, vm.current_locals().and_then(get_locals)?)
220+
(
221+
globals.clone(),
222+
self.locals
223+
.unwrap_or_else(|| PyMapping::new(globals.into())),
224+
)
237225
}
226+
None => (
227+
vm.current_globals().clone(),
228+
if let Some(locals) = self.locals {
229+
locals
230+
} else {
231+
vm.current_locals()?
232+
},
233+
),
238234
};
239235

240236
let scope = Scope::with_builtins(Some(locals), globals, vm);
@@ -432,7 +428,7 @@ mod builtins {
432428
}
433429

434430
#[pyfunction]
435-
fn locals(vm: &VirtualMachine) -> PyResult<PyObjectRef> {
431+
fn locals(vm: &VirtualMachine) -> PyResult<PyMapping> {
436432
vm.current_locals()
437433
}
438434

@@ -799,7 +795,7 @@ mod builtins {
799795
vm.new_type_error("vars() argument must have __dict__ attribute".to_owned())
800796
})
801797
} else {
802-
Ok(vm.current_locals()?)
798+
Ok(vm.current_locals()?.into_object())
803799
}
804800
}
805801

@@ -876,13 +872,22 @@ mod builtins {
876872
FuncArgs::new(vec![name_obj.clone().into(), bases.clone()], kwargs.clone()),
877873
)?;
878874

879-
let namespace = PyDictRef::try_from_object(vm, namespace)?;
875+
// Accept any PyMapping as namespace.
876+
let namespace = PyMapping::try_from_object(vm, namespace.clone()).map_err(|_| {
877+
vm.new_type_error(format!(
878+
"{}.__prepare__() must return a mapping, not {}",
879+
metaclass.name(),
880+
namespace.class().name()
881+
))
882+
})?;
880883

881884
let classcell = function.invoke_with_locals(().into(), Some(namespace.clone()), vm)?;
882885
let classcell = <Option<PyCellRef>>::try_from_object(vm, classcell)?;
883886

884887
if let Some(orig_bases) = orig_bases {
885-
namespace.set_item("__orig_bases__", orig_bases.into(), vm)?;
888+
namespace
889+
.as_object()
890+
.set_item("__orig_bases__", orig_bases.into(), vm)?;
886891
}
887892

888893
let class = vm.invoke(

‎vm/src/vm.rs

+4-4
Original file line numberDiff line numberDiff line change
@@ -22,14 +22,14 @@ use crate::{
2222
frozen,
2323
function::{FuncArgs, IntoFuncArgs, IntoPyObject},
2424
import,
25-
protocol::{PyIterIter, PyIterReturn},
25+
protocol::{PyIterIter, PyIterReturn, PyMapping},
2626
scope::Scope,
2727
signal::NSIG,
2828
stdlib,
2929
types::PyComparisonOp,
3030
utils::Either,
3131
IdProtocol, ItemProtocol, PyArithmeticValue, PyContext, PyLease, PyMethod, PyObject,
32-
PyObjectRef, PyRef, PyRefExact, PyResult, PyValue, TryFromObject, TypeProtocol,
32+
PyObjectRef, PyObjectWrap, PyRef, PyRefExact, PyResult, PyValue, TryFromObject, TypeProtocol,
3333
};
3434
use crossbeam_utils::atomic::AtomicCell;
3535
use num_traits::{Signed, ToPrimitive};
@@ -558,7 +558,7 @@ impl VirtualMachine {
558558
}
559559
}
560560

561-
pub fn current_locals(&self) -> PyResult<PyObjectRef> {
561+
pub fn current_locals(&self) -> PyResult<PyMapping> {
562562
self.current_frame()
563563
.expect("called current_locals but no frames on the stack")
564564
.locals(self)
@@ -1201,7 +1201,7 @@ impl VirtualMachine {
12011201
.get_special_method(obj, "__dir__")?
12021202
.map_err(|_obj| self.new_type_error("object does not provide __dir__".to_owned()))?
12031203
.invoke((), self)?,
1204-
None => self.call_method(&self.current_locals()?, "keys", ())?,
1204+
None => self.call_method(self.current_locals()?.as_object(), "keys", ())?,
12051205
};
12061206
let items = self.extract_elements(&seq)?;
12071207
let lst = PyList::from(items);

0 commit comments

Comments
 (0)