Implements __mul__ for lists

Implements __mul__ for lists so that it is possible to do things like

```
s = [1, 2,] * 3
s == [1, 2, 1, 2, 1, 2]
```
This commit is contained in:
Ross Jones
2018-12-16 11:44:02 +00:00
parent 62c53d8e5d
commit ae2f7ed1cb
2 changed files with 31 additions and 0 deletions

View File

@@ -9,3 +9,5 @@ assert y == [2, 1, 2, 3]
y.extend(x)
assert y == [2, 1, 2, 3, 1, 2, 3]
assert x * 0 == [], "list __mul__ by 0 failed"
assert x * 2 == [1, 2, 3, 1, 2, 3], "list __mul__ by 2 failed"

View File

@@ -232,6 +232,34 @@ fn list_setitem(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
set_item(vm, &mut elements, key.clone(), value.clone())
}
fn list_mul(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
arg_check!(
vm,
args,
required = [
(list, Some(vm.ctx.list_type())),
(product, Some(vm.ctx.int_type()))
]
);
let counter = objint::get_value(&product).to_usize().unwrap();
let elements = get_elements(list);
let current_len = elements.len();
let mut new_elements = Vec::with_capacity(counter * current_len);
for _ in 0..counter {
new_elements.extend(elements.clone());
}
Ok(PyObject::new(
PyObjectKind::Sequence {
elements: new_elements,
},
vm.ctx.list_type(),
))
}
pub fn init(context: &PyContext) {
let ref list_type = context.list_type;
context.set_attr(&list_type, "__add__", context.new_rustfunc(list_add));
@@ -252,6 +280,7 @@ pub fn init(context: &PyContext) {
"__setitem__",
context.new_rustfunc(list_setitem),
);
context.set_attr(&list_type, "__mul__", context.new_rustfunc(list_mul));
context.set_attr(&list_type, "__len__", context.new_rustfunc(list_len));
context.set_attr(&list_type, "__new__", context.new_rustfunc(list_new));
context.set_attr(&list_type, "__repr__", context.new_rustfunc(list_repr));