Skip to content

Commit 525807c

Browse files
windelbouwmanwindelbouwman
committed
Add __getitem__ on list, tuple and str
1 parent f59f8a9 commit 525807c

5 files changed

Lines changed: 42 additions & 14 deletions

File tree

‎tests/snippets/function_args.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,3 +24,8 @@ def va(a, b=2, *c, d, **e):
2424
assert e['f'] == 42
2525

2626
va(1, 22, 3, 4, d=1337, f=42)
27+
28+
def va2(*args, **kwargs):
29+
assert args == (5, 4)
30+
31+
va2(5, 4)

‎vm/src/obj/objlist.rs‎

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ use super::super::pyobject::{
44
use super::super::vm::VirtualMachine;
55
use super::objbool;
66
use super::objiter;
7-
use super::objsequence::{seq_equal, PySliceableSequence};
7+
use super::objsequence::{get_item, seq_equal, PySliceableSequence};
88
use super::objstr;
99
use super::objtype;
1010

@@ -132,7 +132,7 @@ pub fn list_append(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
132132
}
133133
}
134134

135-
fn clear(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
135+
fn list_clear(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
136136
trace!("list.clear called with: {:?}", args);
137137
arg_check!(vm, args, required = [(list, Some(vm.ctx.list_type()))]);
138138
let mut list_obj = list.borrow_mut();
@@ -184,15 +184,26 @@ fn list_contains(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
184184
Ok(vm.new_bool(false))
185185
}
186186

187+
fn list_getitem(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
188+
trace!("list.getitem called with: {:?}", args);
189+
arg_check!(
190+
vm,
191+
args,
192+
required = [(list, Some(vm.ctx.list_type())), (needle, None)]
193+
);
194+
get_item(vm, list, &get_elements(list), needle.clone())
195+
}
196+
187197
pub fn init(context: &PyContext) {
188198
let ref list_type = context.list_type;
189199
list_type.set_attr("__add__", context.new_rustfunc(list_add));
190200
list_type.set_attr("__contains__", context.new_rustfunc(list_contains));
191201
list_type.set_attr("__eq__", context.new_rustfunc(list_eq));
202+
list_type.set_attr("__getitem__", context.new_rustfunc(list_getitem));
192203
list_type.set_attr("__len__", context.new_rustfunc(list_len));
193204
list_type.set_attr("__new__", context.new_rustfunc(list_new));
194205
list_type.set_attr("__repr__", context.new_rustfunc(list_repr));
195206
list_type.set_attr("append", context.new_rustfunc(list_append));
196-
list_type.set_attr("clear", context.new_rustfunc(clear));
207+
list_type.set_attr("clear", context.new_rustfunc(list_clear));
197208
list_type.set_attr("reverse", context.new_rustfunc(list_reverse));
198209
}

‎vm/src/obj/objstr.rs‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ pub fn init(context: &PyContext) {
1111
str_type.set_attr("__add__", context.new_rustfunc(str_add));
1212
str_type.set_attr("__eq__", context.new_rustfunc(str_eq));
1313
str_type.set_attr("__contains__", context.new_rustfunc(str_contains));
14+
str_type.set_attr("__getitem__", context.new_rustfunc(str_getitem));
1415
str_type.set_attr("__len__", context.new_rustfunc(str_len));
1516
str_type.set_attr("__mul__", context.new_rustfunc(str_mul));
1617
str_type.set_attr("__new__", context.new_rustfunc(str_new));
@@ -221,6 +222,16 @@ fn str_contains(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
221222
Ok(vm.ctx.new_bool(value.contains(needle.as_str())))
222223
}
223224

225+
fn str_getitem(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
226+
arg_check!(
227+
vm,
228+
args,
229+
required = [(s, Some(vm.ctx.str_type())), (needle, None)]
230+
);
231+
let value = get_value(&s);
232+
subscript(vm, &value, needle.clone())
233+
}
234+
224235
// TODO: should with following format
225236
// class str(object='')
226237
// class str(object=b'', encoding='utf-8', errors='strict')

‎vm/src/obj/objtuple.rs‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ use super::super::pyobject::{
33
};
44
use super::super::vm::VirtualMachine;
55
use super::objbool;
6-
use super::objsequence::seq_equal;
6+
use super::objsequence::{get_item, seq_equal};
77
use super::objstr;
88
use super::objtype;
99

@@ -51,6 +51,15 @@ fn tuple_repr(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
5151
Ok(vm.new_str(s))
5252
}
5353

54+
fn tuple_getitem(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
55+
arg_check!(
56+
vm,
57+
args,
58+
required = [(tuple, Some(vm.ctx.tuple_type())), (needle, None)]
59+
);
60+
get_item(vm, tuple, &get_elements(&tuple), needle.clone())
61+
}
62+
5463
pub fn tuple_contains(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
5564
arg_check!(
5665
vm,
@@ -83,6 +92,7 @@ pub fn init(context: &PyContext) {
8392
let ref tuple_type = context.tuple_type;
8493
tuple_type.set_attr("__eq__", context.new_rustfunc(tuple_eq));
8594
tuple_type.set_attr("__contains__", context.new_rustfunc(tuple_contains));
95+
tuple_type.set_attr("__getitem__", context.new_rustfunc(tuple_getitem));
8696
tuple_type.set_attr("__len__", context.new_rustfunc(tuple_len));
8797
tuple_type.set_attr("__repr__", context.new_rustfunc(tuple_repr));
8898
}

‎vm/src/vm.rs‎

Lines changed: 1 addition & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -411,16 +411,7 @@ impl VirtualMachine {
411411
}
412412

413413
fn subscript(&mut self, a: PyObjectRef, b: PyObjectRef) -> PyResult {
414-
// Subscript implementation: a[b]
415-
// TODO: simply always call __getitem__
416-
let a2 = &*a.borrow();
417-
match &a2.kind {
418-
PyObjectKind::String { ref value } => objstr::subscript(self, value, b),
419-
PyObjectKind::List { ref elements } | PyObjectKind::Tuple { ref elements } => {
420-
super::obj::objsequence::get_item(self, &a, elements, b)
421-
}
422-
_ => self.call_method(&a, "__getitem__", vec![b]),
423-
}
414+
self.call_method(&a, "__getitem__", vec![b])
424415
}
425416

426417
fn execute_store_subscript(&mut self) -> Option<PyResult> {

0 commit comments

Comments
 (0)