Support index in list.pop()

Author: adrian17

tests/snippets/list.py

@@ -29,8 +29,22 @@
 assert x > y, "list __gt__ failed"
 
 
-assert [1,2,'a'].pop() == 'a', "list pop failed"
+x = [0, 1, 2]
+assert x.pop() == 2
+assert x == [0, 1]
+
+def test_pop(lst, idx, value, new_lst):
+    assert lst.pop(idx) == value
+    assert lst == new_lst
+test_pop([0, 1, 2], -1, 2, [0, 1])
+test_pop([0, 1, 2], 0, 0, [1, 2])
+test_pop([0, 1, 2], 1, 1, [0, 2])
+test_pop([0, 1, 2], 2, 2, [0, 1])
 assert_raises(IndexError, lambda: [].pop())
+assert_raises(IndexError, lambda: [].pop(0))
+assert_raises(IndexError, lambda: [].pop(-1))
+assert_raises(IndexError, lambda: [0].pop(1))
+assert_raises(IndexError, lambda: [0].pop(-2))
 
 recursive = []
 recursive.append(recursive)

vm/src/obj/objlist.rs

@@ -203,12 +203,18 @@ impl PyListRef {
         Err(vm.new_value_error(format!("'{}' is not in list", needle_str)))
     }
 
-    fn pop(self, vm: &mut VirtualMachine) -> PyResult {
+    fn pop(self, i: OptionalArg<isize>, vm: &mut VirtualMachine) -> PyResult {
+        let mut i = i.into_option().unwrap_or(-1);
         let mut elements = self.elements.borrow_mut();
-        if let Some(result) = elements.pop() {
-            Ok(result)
-        } else {
+        if i < 0 {
+            i += elements.len() as isize;
+        }
+        if elements.is_empty() {
             Err(vm.new_index_error("pop from empty list".to_string()))
+        } else if i < 0 || i as usize >= elements.len() {
+            Err(vm.new_index_error("pop index out of range".to_string()))
+        } else {
+            Ok(elements.remove(i as usize))
         }
     }