Skip to content

Commit 45290b1

Browse files
authored
Merge pull request #3314 from DimitrisJim/mapping_locals
Allow any mapping for locals.
2 parents 7233579 + 94b9163 commit 45290b1

File tree

8 files changed

+148
-69
lines changed

8 files changed

+148
-69
lines changed

Lib/test/test_builtin.py

-2
Original file line numberDiff line numberDiff line change
@@ -490,8 +490,6 @@ def __getitem__(self, key):
490490
raise ValueError
491491
self.assertRaises(ValueError, eval, "foo", {}, X())
492492

493-
# TODO: RUSTPYTHON
494-
@unittest.expectedFailure
495493
def test_general_eval(self):
496494
# Tests that general mappings can be used for the locals argument
497495

Lib/test/test_compile.py

-2
Original file line numberDiff line numberDiff line change
@@ -61,8 +61,6 @@ def test_none_keyword_arg(self):
6161
def test_duplicate_global_local(self):
6262
self.assertRaises(SyntaxError, exec, 'def f(a): global a; a = 1')
6363

64-
# TODO: RUSTPYTHON
65-
@unittest.expectedFailure
6664
def test_exec_with_general_mapping_for_locals(self):
6765

6866
class M:

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<PyDictRef> {
48+
fn f_locals(self, vm: &VirtualMachine) -> PyResult<PyMapping> {
4849
self.locals(vm)
4950
}
5051

vm/src/builtins/function.rs

+6-3
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,9 +288,11 @@ 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())
295+
PyMapping::new(self.globals.clone().into())
293296
};
294297

295298
// Construct frame:

vm/src/frame.rs

+102-35
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: PyDictRef,
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<PyDictRef> {
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.as_object().clone().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 PyDictRef,
292+
locals: &'a PyMapping,
279293
globals: &'a PyDictRef,
280294
builtins: &'a PyDictRef,
281295
object: &'a FrameRef,
@@ -496,12 +510,20 @@ impl ExecutingFrame<'_> {
496510
}
497511
bytecode::Instruction::LoadNameAny(idx) => {
498512
let name = &self.code.names[*idx as usize];
499-
let x = self.locals.get_item_option(name.clone(), vm)?;
500-
let x = match x {
513+
// Try using locals as dict first, if not, fallback to generic method.
514+
let x = match self
515+
.locals
516+
.clone()
517+
.into_object()
518+
.downcast_exact::<PyDict>(vm)
519+
{
520+
Ok(d) => d.get_item_option(name.clone(), vm)?,
521+
Err(o) => o.get_item(name.clone(), vm).ok(),
522+
};
523+
self.push_value(match x {
501524
Some(x) => x,
502525
None => self.load_global_or_builtin(name, vm)?,
503-
};
504-
self.push_value(x);
526+
});
505527
Ok(None)
506528
}
507529
bytecode::Instruction::LoadGlobal(idx) => {
@@ -521,14 +543,22 @@ impl ExecutingFrame<'_> {
521543
bytecode::Instruction::LoadClassDeref(i) => {
522544
let i = *i as usize;
523545
let name = self.code.freevars[i - self.code.cellvars.len()].clone();
524-
let value = if let Some(value) = self.locals.get_item_option(name, vm)? {
525-
value
526-
} else {
527-
self.cells_frees[i]
528-
.get()
529-
.ok_or_else(|| self.unbound_cell_exception(i, vm))?
546+
// Try using locals as dict first, if not, fallback to generic method.
547+
let value = match self
548+
.locals
549+
.clone()
550+
.into_object()
551+
.downcast_exact::<PyDict>(vm)
552+
{
553+
Ok(d) => d.get_item_option(name, vm)?,
554+
Err(o) => o.get_item(name, vm).ok(),
530555
};
531-
self.push_value(value);
556+
self.push_value(match value {
557+
Some(v) => v,
558+
None => self.cells_frees[i]
559+
.get()
560+
.ok_or_else(|| self.unbound_cell_exception(i, vm))?,
561+
});
532562
Ok(None)
533563
}
534564
bytecode::Instruction::StoreFast(idx) => {
@@ -538,8 +568,15 @@ impl ExecutingFrame<'_> {
538568
}
539569
bytecode::Instruction::StoreLocal(idx) => {
540570
let value = self.pop_value();
541-
self.locals
542-
.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+
};
543580
Ok(None)
544581
}
545582
bytecode::Instruction::StoreGlobal(idx) => {
@@ -559,7 +596,17 @@ impl ExecutingFrame<'_> {
559596
}
560597
bytecode::Instruction::DeleteLocal(idx) => {
561598
let name = &self.code.names[*idx as usize];
562-
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 {
563610
Ok(()) => {}
564611
Err(e) if e.isinstance(&vm.ctx.exceptions.key_error) => {
565612
return Err(vm.new_name_error(format!("name '{}' is not defined", name)))
@@ -717,9 +764,25 @@ impl ExecutingFrame<'_> {
717764
}
718765
bytecode::Instruction::YieldFrom => self.execute_yield_from(vm),
719766
bytecode::Instruction::SetupAnnotation => {
720-
if !self.locals.contains_key("__annotations__", vm) {
721-
self.locals
722-
.set_item("__annotations__", vm.ctx.new_dict().into(), vm)?;
767+
// Try using locals as dict first, if not, fallback to generic method.
768+
let has_annotations = match self
769+
.locals
770+
.clone()
771+
.into_object()
772+
.downcast_exact::<PyDict>(vm)
773+
{
774+
Ok(d) => d.contains_key("__annotations__", vm),
775+
Err(o) => {
776+
let needle = vm.new_pyobj("__annotations__");
777+
self._in(vm, needle, o)?
778+
}
779+
};
780+
if !has_annotations {
781+
self.locals.as_object().set_item(
782+
"__annotations__",
783+
vm.ctx.new_dict().into(),
784+
vm,
785+
)?;
723786
}
724787
Ok(None)
725788
}
@@ -1132,7 +1195,15 @@ impl ExecutingFrame<'_> {
11321195
for (k, v) in &dict {
11331196
let k = PyStrRef::try_from_object(vm, k)?;
11341197
if filter_pred(k.as_str()) {
1135-
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+
};
11361207
}
11371208
}
11381209
}
@@ -1822,15 +1893,11 @@ impl fmt::Debug for Frame {
18221893
.map(|elem| format!("\n > {:?}", elem))
18231894
.collect::<String>();
18241895
// TODO: fix this up
1825-
let dict = self.locals.clone();
1826-
let local_str = dict
1827-
.into_iter()
1828-
.map(|elem| format!("\n {:?} = {:?}", elem.0, elem.1))
1829-
.collect::<String>();
1896+
let locals = self.locals.clone();
18301897
write!(
18311898
f,
1832-
"Frame Object {{ \n Stack:{}\n Blocks:{}\n Locals:{}\n}}",
1833-
stack_str, block_str, local_str
1899+
"Frame Object {{ \n Stack:{}\n Blocks:{}\n Locals:{:?}\n}}",
1900+
stack_str, block_str, locals
18341901
)
18351902
}
18361903
}

vm/src/scope.rs

+5-4
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+
protocol::PyMapping,
45
ItemProtocol, VirtualMachine,
56
};
67
use std::fmt;
78

89
#[derive(Clone)]
910
pub struct Scope {
10-
pub locals: PyDictRef,
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<PyDictRef>, globals: PyDictRef) -> Scope {
24-
let locals = locals.unwrap_or_else(|| globals.clone());
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<PyDictRef>,
30+
locals: Option<PyMapping>,
3031
globals: PyDictRef,
3132
vm: &VirtualMachine,
3233
) -> Scope {

0 commit comments

Comments
 (0)