Implements __mul__ for lists

rossjones · rossjones · commit ae2f7ed1cbca · 2018-12-16T12:22:31.000Z
Implements __mul__ for lists so that it is possible to do things like

```
s = [1, 2,] * 3
s == [1, 2, 1, 2, 1, 2]
```
diff --git a/tests/snippets/list.py b/tests/snippets/list.py
@@ -9,3 +9,5 @@
 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"
diff --git a/vm/src/obj/objlist.rs b/vm/src/obj/objlist.rs
@@ -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));
