Skip to content

Commit b123e58

Browse files
committed
Support index in list.pop()
1 parent 8143adc commit b123e58

File tree

2 files changed

+25
-5
lines changed

2 files changed

+25
-5
lines changed

tests/snippets/list.py

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,8 +29,22 @@
2929
assert x > y, "list __gt__ failed"
3030

3131

32-
assert [1,2,'a'].pop() == 'a', "list pop failed"
32+
x = [0, 1, 2]
33+
assert x.pop() == 2
34+
assert x == [0, 1]
35+
36+
def test_pop(lst, idx, value, new_lst):
37+
assert lst.pop(idx) == value
38+
assert lst == new_lst
39+
test_pop([0, 1, 2], -1, 2, [0, 1])
40+
test_pop([0, 1, 2], 0, 0, [1, 2])
41+
test_pop([0, 1, 2], 1, 1, [0, 2])
42+
test_pop([0, 1, 2], 2, 2, [0, 1])
3343
assert_raises(IndexError, lambda: [].pop())
44+
assert_raises(IndexError, lambda: [].pop(0))
45+
assert_raises(IndexError, lambda: [].pop(-1))
46+
assert_raises(IndexError, lambda: [0].pop(1))
47+
assert_raises(IndexError, lambda: [0].pop(-2))
3448

3549
recursive = []
3650
recursive.append(recursive)

vm/src/obj/objlist.rs

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -203,12 +203,18 @@ impl PyListRef {
203203
Err(vm.new_value_error(format!("'{}' is not in list", needle_str)))
204204
}
205205

206-
fn pop(self, vm: &mut VirtualMachine) -> PyResult {
206+
fn pop(self, i: OptionalArg<isize>, vm: &mut VirtualMachine) -> PyResult {
207+
let mut i = i.into_option().unwrap_or(-1);
207208
let mut elements = self.elements.borrow_mut();
208-
if let Some(result) = elements.pop() {
209-
Ok(result)
210-
} else {
209+
if i < 0 {
210+
i += elements.len() as isize;
211+
}
212+
if elements.is_empty() {
211213
Err(vm.new_index_error("pop from empty list".to_string()))
214+
} else if i < 0 || i as usize >= elements.len() {
215+
Err(vm.new_index_error("pop index out of range".to_string()))
216+
} else {
217+
Ok(elements.remove(i as usize))
212218
}
213219
}
214220

0 commit comments

Comments
 (0)