Skip to content

Commit 895b52e

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

File tree

6 files changed

+124
-58
lines changed

6 files changed

+124
-58
lines changed

‎vm/src/builtins/frame.rs

Lines changed: 2 additions & 1 deletion
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

Lines changed: 7 additions & 5 deletions
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

Lines changed: 76 additions & 19 deletions
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,16 @@ 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+
match locals.as_object().clone().downcast_exact::<PyDict>(vm) {
192+
Ok(d) => d.set_item(k.clone(), v.clone(), vm)?,
193+
Err(o) => o.set_item(k.clone(), v.clone(), vm)?,
194+
};
192195
} else {
193-
match locals.del_item(k.clone(), vm) {
196+
let res = match locals.as_object().clone().downcast_exact::<PyDict>(vm) {
197+
Ok(d) => d.del_item(k.clone(), vm),
198+
Err(o) => o.del_item(k.clone(), vm),
199+
};
200+
match res {
194201
Ok(()) => {}
195202
Err(e) if e.isinstance(&vm.ctx.exceptions.key_error) => {}
196203
Err(e) => return Err(e),
@@ -202,9 +209,16 @@ impl FrameRef {
202209
let map_to_dict = |keys: &[PyStrRef], values: &[PyCellRef]| {
203210
for (k, v) in itertools::zip(keys, values) {
204211
if let Some(v) = v.get() {
205-
locals.set_item(k.clone(), v, vm)?;
212+
match locals.into_object().downcast_exact::<PyDict>(vm) {
213+
Ok(d) => d.set_item(k.clone(), v, vm)?,
214+
Err(o) => o.set_item(k.clone(), v, vm)?,
215+
};
206216
} else {
207-
match locals.del_item(k.clone(), vm) {
217+
let res = match locals.as_object().clone().downcast_exact::<PyDict>(vm) {
218+
Ok(d) => d.del_item(k.clone(), vm),
219+
Err(o) => o.del_item(k.clone(), vm),
220+
};
221+
match res {
208222
Ok(()) => {}
209223
Err(e) if e.isinstance(&vm.ctx.exceptions.key_error) => {}
210224
Err(e) => return Err(e),
@@ -275,7 +289,7 @@ struct ExecutingFrame<'a> {
275289
code: &'a PyRef<PyCode>,
276290
fastlocals: &'a PyMutex<Box<[Option<PyObjectRef>]>>,
277291
cells_frees: &'a [PyCellRef],
278-
locals: &'a PyObjectRef,
292+
locals: &'a PyMapping,
279293
globals: &'a PyDictRef,
280294
builtins: &'a PyDictRef,
281295
object: &'a FrameRef,
@@ -497,7 +511,12 @@ impl ExecutingFrame<'_> {
497511
bytecode::Instruction::LoadNameAny(idx) => {
498512
let name = &self.code.names[*idx as usize];
499513
// Try using locals as dict first, if not, fallback to generic method.
500-
let x = match self.locals.clone().downcast_exact::<PyDict>(vm) {
514+
let x = match self
515+
.locals
516+
.clone()
517+
.into_object()
518+
.downcast_exact::<PyDict>(vm)
519+
{
501520
Ok(d) => d.get_item_option(name.clone(), vm)?,
502521
Err(o) => o.get_item(name.clone(), vm).ok(),
503522
};
@@ -525,7 +544,12 @@ impl ExecutingFrame<'_> {
525544
let i = *i as usize;
526545
let name = self.code.freevars[i - self.code.cellvars.len()].clone();
527546
// Try using locals as dict first, if not, fallback to generic method.
528-
let value = match self.locals.clone().downcast_exact::<PyDict>(vm) {
547+
let value = match self
548+
.locals
549+
.clone()
550+
.into_object()
551+
.downcast_exact::<PyDict>(vm)
552+
{
529553
Ok(d) => d.get_item_option(name, vm)?,
530554
Err(o) => o.get_item(name, vm).ok(),
531555
};
@@ -544,8 +568,15 @@ impl ExecutingFrame<'_> {
544568
}
545569
bytecode::Instruction::StoreLocal(idx) => {
546570
let value = self.pop_value();
547-
self.locals
548-
.set_item(self.code.names[*idx as usize].clone(), value, vm)?;
571+
match self
572+
.locals
573+
.clone()
574+
.into_object()
575+
.downcast_exact::<PyDict>(vm)
576+
{
577+
Ok(d) => d.set_item(self.code.names[*idx as usize].clone(), value, vm)?,
578+
Err(o) => o.set_item(self.code.names[*idx as usize].clone(), value, vm)?,
579+
};
549580
Ok(None)
550581
}
551582
bytecode::Instruction::StoreGlobal(idx) => {
@@ -565,7 +596,17 @@ impl ExecutingFrame<'_> {
565596
}
566597
bytecode::Instruction::DeleteLocal(idx) => {
567598
let name = &self.code.names[*idx as usize];
568-
match self.locals.del_item(name.clone(), vm) {
599+
let res = match self
600+
.locals
601+
.clone()
602+
.into_object()
603+
.downcast_exact::<PyDict>(vm)
604+
{
605+
Ok(d) => d.del_item(name.clone(), vm),
606+
Err(o) => o.del_item(name.clone(), vm),
607+
};
608+
609+
match res {
569610
Ok(()) => {}
570611
Err(e) if e.isinstance(&vm.ctx.exceptions.key_error) => {
571612
return Err(vm.new_name_error(format!("name '{}' is not defined", name)))
@@ -724,16 +765,24 @@ impl ExecutingFrame<'_> {
724765
bytecode::Instruction::YieldFrom => self.execute_yield_from(vm),
725766
bytecode::Instruction::SetupAnnotation => {
726767
// Try using locals as dict first, if not, fallback to generic method.
727-
let has_annotations = match self.locals.clone().downcast_exact::<PyDict>(vm) {
768+
let has_annotations = match self
769+
.locals
770+
.clone()
771+
.into_object()
772+
.downcast_exact::<PyDict>(vm)
773+
{
728774
Ok(d) => d.contains_key("__annotations__", vm),
729775
Err(o) => {
730776
let needle = vm.new_pyobj("__annotations__");
731777
self._in(vm, needle, o)?
732778
}
733779
};
734780
if !has_annotations {
735-
self.locals
736-
.set_item("__annotations__", vm.ctx.new_dict().into(), vm)?;
781+
self.locals.as_object().set_item(
782+
"__annotations__",
783+
vm.ctx.new_dict().into(),
784+
vm,
785+
)?;
737786
}
738787
Ok(None)
739788
}
@@ -1146,7 +1195,15 @@ impl ExecutingFrame<'_> {
11461195
for (k, v) in &dict {
11471196
let k = PyStrRef::try_from_object(vm, k)?;
11481197
if filter_pred(k.as_str()) {
1149-
self.locals.set_item(k, v, vm)?;
1198+
match self
1199+
.locals
1200+
.clone()
1201+
.into_object()
1202+
.downcast_exact::<PyDict>(vm)
1203+
{
1204+
Ok(d) => d.set_item(k, v, vm)?,
1205+
Err(o) => o.set_item(k, v, vm)?,
1206+
};
11501207
}
11511208
}
11521209
}

‎vm/src/scope.rs

Lines changed: 6 additions & 5 deletions
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

Lines changed: 29 additions & 24 deletions
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(

0 commit comments

Comments
 (0)