Merge pull request RustPython#503 from janczer/add_pop_bytearray

Add pop to bytearray

Author: OddCoincidence

tests/snippets/bytearray.py

@@ -33,3 +33,22 @@
 a = bytearray(b'abcd')
 a.clear()
 assert len(a) == 0
+
+b = bytearray(b'test')
+assert len(b) == 4
+b.pop()
+assert len(b) == 3
+
+c = bytearray([123, 255, 111])
+assert len(c) == 3
+c.pop()
+assert len(c) == 2
+c.pop()
+c.pop()
+
+try:
+    c.pop()
+except IndexError:
+    pass
+else:
+    assert False

vm/src/obj/objbytearray.rs

@@ -7,6 +7,7 @@ use super::super::pyobject::{
 use super::objint;
 
 use super::super::vm::VirtualMachine;
+use super::objbytes::get_mut_value;
 use super::objbytes::get_value;
 use super::objtype;
 use num_traits::ToPrimitive;
@@ -100,6 +101,7 @@ pub fn init(context: &PyContext) {
         "clear",
         context.new_rustfunc(bytearray_clear),
     );
+    context.set_attr(&bytearray_type, "pop", context.new_rustfunc(bytearray_pop));
 }
 
 fn bytearray_new(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
@@ -278,3 +280,14 @@ fn bytearray_clear(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
         _ => panic!("Bytearray has incorrect payload."),
     }
 }
+
+fn bytearray_pop(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+    arg_check!(vm, args, required = [(obj, Some(vm.ctx.bytearray_type()))]);
+    let mut value = get_mut_value(obj);
+
+    if let Some(i) = value.pop() {
+        Ok(vm.ctx.new_int(i))
+    } else {
+        Err(vm.new_index_error("pop from empty bytearray".to_string()))
+    }
+}

vm/src/obj/objbytes.rs

@@ -6,8 +6,10 @@ use super::objint;
 use super::objtype;
 use num_traits::ToPrimitive;
 use std::cell::Ref;
+use std::cell::RefMut;
 use std::hash::{Hash, Hasher};
 use std::ops::Deref;
+use std::ops::DerefMut;
 
 // Binary data support
 
@@ -183,7 +185,17 @@ pub fn get_value<'a>(obj: &'a PyObjectRef) -> impl Deref<Target = Vec<u8>> + 'a
         if let PyObjectPayload::Bytes { ref value } = py_obj.payload {
             value
         } else {
-            panic!("Inner error getting int {:?}", obj);
+            panic!("Inner error getting bytearray {:?}", obj);
+        }
+    })
+}
+
+pub fn get_mut_value<'a>(obj: &'a PyObjectRef) -> impl DerefMut<Target = Vec<u8>> + 'a {
+    RefMut::map(obj.borrow_mut(), |py_obj| {
+        if let PyObjectPayload::Bytes { ref mut value } = py_obj.payload {
+            value
+        } else {
+            panic!("Inner error getting bytearray {:?}", obj);
         }
     })
 }