diff --git a/tests/snippets/list.py b/tests/snippets/list.py index 9e385abdce..6b9839321d 100644 --- a/tests/snippets/list.py +++ b/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) diff --git a/vm/src/obj/objlist.rs b/vm/src/obj/objlist.rs index 8ad69bff31..477ebcf5e4 100644 --- a/vm/src/obj/objlist.rs +++ b/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, 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)) } }