Add list.extend method

Author: windelbouwman

tests/snippets/list.py

@@ -6,3 +6,6 @@
 y = [2, *x]
 assert y == [2, 1, 2, 3]
 
+y.extend(x)
+assert y == [2, 1, 2, 3, 1, 2, 3]
+

vm/src/frame.rs

@@ -13,7 +13,6 @@ use super::obj::objdict;
 use super::obj::objiter;
 use super::obj::objlist;
 use super::obj::objstr;
-use super::obj::objtuple;
 use super::obj::objtype;
 use super::pyobject::{
     AttributeProtocol, DictProtocol, IdProtocol, ParentProtocol, PyFuncArgs, PyObject,
@@ -482,7 +481,7 @@ impl Frame {
                             vec![]
                         };
                         let args = self.pop_value();
-                        let args = self.extract_elements(vm, &args)?;
+                        let args = vm.extract_elements(&args)?;
                         PyFuncArgs { args, kwargs }
                     }
                 };
@@ -595,7 +594,7 @@ impl Frame {
             }
             bytecode::Instruction::UnpackSequence { size } => {
                 let value = self.pop_value();
-                let elements = self.extract_elements(vm, &value)?;
+                let elements = vm.extract_elements(&value)?;
                 if elements.len() != *size {
                     Err(vm.new_value_error("Wrong number of values to unpack".to_string()))
                 } else {
@@ -607,7 +606,7 @@ impl Frame {
             }
             bytecode::Instruction::UnpackEx { before, after } => {
                 let value = self.pop_value();
-                let elements = self.extract_elements(vm, &value)?;
+                let elements = vm.extract_elements(&value)?;
                 let min_expected = *before + *after;
                 if elements.len() < min_expected {
                     Err(vm.new_value_error(format!(
@@ -642,7 +641,7 @@ impl Frame {
             }
             bytecode::Instruction::Unpack => {
                 let value = self.pop_value();
-                let elements = self.extract_elements(vm, &value)?;
+                let elements = vm.extract_elements(&value)?;
                 for element in elements.into_iter().rev() {
                     self.push_value(element);
                 }
@@ -651,23 +650,6 @@ impl Frame {
         }
     }
 
-    fn extract_elements(
-        &mut self,
-        vm: &mut VirtualMachine,
-        value: &PyObjectRef,
-    ) -> Result<Vec<PyObjectRef>, PyObjectRef> {
-        // Extract elements from item, if possible:
-        let elements = if objtype::isinstance(value, &vm.ctx.tuple_type()) {
-            objtuple::get_elements(value)
-        } else if objtype::isinstance(value, &vm.ctx.list_type()) {
-            objlist::get_elements(value)
-        } else {
-            let iter = objiter::get_iter(vm, value)?;
-            objiter::get_all(vm, &iter)?
-        };
-        Ok(elements)
-    }
-
     fn get_elements(
         &mut self,
         vm: &mut VirtualMachine,
@@ -678,7 +660,7 @@ impl Frame {
         if unpack {
             let mut result: Vec<PyObjectRef> = vec![];
             for element in elements {
-                let expanded = self.extract_elements(vm, &element)?;
+                let expanded = vm.extract_elements(&element)?;
                 for inner in expanded {
                     result.push(inner);
                 }

vm/src/obj/objlist.rs

@@ -137,6 +137,22 @@ fn list_clear(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     }
 }
 
+pub fn list_extend(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+    arg_check!(
+        vm,
+        args,
+        required = [(list, Some(vm.ctx.list_type())), (x, None)]
+    );
+    let mut new_elements = vm.extract_elements(x)?;
+    let mut list_obj = list.borrow_mut();
+    if let PyObjectKind::List { ref mut elements } = list_obj.kind {
+        elements.append(&mut new_elements);
+        Ok(vm.get_none())
+    } else {
+        Err(vm.new_type_error("list.extend is called with no list".to_string()))
+    }
+}
+
 fn list_len(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     trace!("list.len called with: {:?}", args);
     arg_check!(vm, args, required = [(list, Some(vm.ctx.list_type()))]);
@@ -198,5 +214,6 @@ pub fn init(context: &PyContext) {
     list_type.set_attr("__repr__", context.new_rustfunc(list_repr));
     list_type.set_attr("append", context.new_rustfunc(list_append));
     list_type.set_attr("clear", context.new_rustfunc(list_clear));
+    list_type.set_attr("extend", context.new_rustfunc(list_extend));
     list_type.set_attr("reverse", context.new_rustfunc(list_reverse));
 }

vm/src/vm.rs

@@ -12,7 +12,10 @@ use super::builtins;
 use super::bytecode;
 use super::frame::{copy_code, Frame};
 use super::obj::objgenerator;
+use super::obj::objiter;
+use super::obj::objlist;
 use super::obj::objobject;
+use super::obj::objtuple;
 use super::obj::objtype;
 use super::pyobject::{DictProtocol, PyContext, PyFuncArgs, PyObjectKind, PyObjectRef, PyResult};
 use super::stdlib;
@@ -351,6 +354,22 @@ impl VirtualMachine {
         Ok(())
     }
 
+    pub fn extract_elements(
+        &mut self,
+        value: &PyObjectRef,
+    ) -> Result<Vec<PyObjectRef>, PyObjectRef> {
+        // Extract elements from item, if possible:
+        let elements = if objtype::isinstance(value, &self.ctx.tuple_type()) {
+            objtuple::get_elements(value)
+        } else if objtype::isinstance(value, &self.ctx.list_type()) {
+            objlist::get_elements(value)
+        } else {
+            let iter = objiter::get_iter(self, value)?;
+            objiter::get_all(self, &iter)?
+        };
+        Ok(elements)
+    }
+
     pub fn get_attribute(&mut self, obj: PyObjectRef, attr_name: &str) -> PyResult {
         objtype::get_attribute(self, obj.clone(), attr_name)
     }