Add bytearray object

windelbouwman · windelbouwman · commit 810ef7df1c52 · 2018-11-03T11:09:15.000+01:00
diff --git a/tests/snippets/whats_left_to_implement.py b/tests/snippets/whats_left_to_implement.py
@@ -858,12 +858,11 @@
         not_implemented.append(("bool", method))
 
 for method in bytearray_expected_methods:
-# TODO: uncomment this when bytearray is implemented
-#    try:
-#        if not hasattr(bytearray(), method):
-#            not_implemented.append(("bytearray", method))
-#    except NameError:
-    not_implemented.append(("bytearray", method))
+    try:
+        if not hasattr(bytearray(), method):
+            not_implemented.append(("bytearray", method))
+    except NameError:
+        not_implemented.append(("bytearray", method))
 
 for method in bytes_expected_methods:
     try:
diff --git a/vm/src/builtins.rs b/vm/src/builtins.rs
@@ -96,7 +96,6 @@ fn builtin_bin(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
 }
 
 // builtin_breakpoint
-// builtin_bytearray
 // builtin_callable
 
 fn builtin_chr(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
@@ -531,6 +530,7 @@ pub fn make_module(ctx: &PyContext) -> PyObjectRef {
     dict.insert(String::from("any"), ctx.new_rustfunc(builtin_any));
     dict.insert(String::from("bin"), ctx.new_rustfunc(builtin_bin));
     dict.insert(String::from("bool"), ctx.bool_type());
+    dict.insert(String::from("bytearray"), ctx.bytearray_type());
     dict.insert(String::from("bytes"), ctx.bytes_type());
     dict.insert(String::from("chr"), ctx.new_rustfunc(builtin_chr));
     dict.insert(String::from("compile"), ctx.new_rustfunc(builtin_compile));
diff --git a/vm/src/obj/mod.rs b/vm/src/obj/mod.rs
@@ -1,6 +1,7 @@
 //! This package contains the python basic/builtin types
 
 pub mod objbool;
+pub mod objbytearray;
 pub mod objbytes;
 pub mod objcomplex;
 pub mod objdict;
diff --git a/vm/src/obj/objbytearray.rs b/vm/src/obj/objbytearray.rs
@@ -0,0 +1,80 @@
+//! Implementation of the python bytearray object.
+
+use super::super::pyobject::{
+    AttributeProtocol, PyContext, PyFuncArgs, PyObject, PyObjectKind, PyObjectRef, PyResult,
+    TypeProtocol,
+};
+use super::super::vm::VirtualMachine;
+use super::objbytes::get_value;
+use super::objint;
+use super::objtype;
+use num_traits::ToPrimitive;
+// Binary data support
+
+/// Fill bytearray class methods dictionary.
+pub fn init(context: &PyContext) {
+    let ref bytearray_type = context.bytearray_type;
+    bytearray_type.set_attr("__eq__", context.new_rustfunc(bytearray_eq));
+    bytearray_type.set_attr("__new__", context.new_rustfunc(bytearray_new));
+    bytearray_type.set_attr("__repr__", context.new_rustfunc(bytearray_repr));
+}
+
+fn bytearray_new(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+    arg_check!(
+        vm,
+        args,
+        required = [(cls, None)],
+        optional = [(val_option, None)]
+    );
+    if !objtype::issubclass(cls, &vm.ctx.bytearray_type()) {
+        return Err(vm.new_type_error(format!("{:?} is not a subtype of bytearray", cls)));
+    }
+
+    // Create bytes data:
+    let value = if let Some(ival) = val_option {
+        let elements = vm.extract_elements(ival)?;
+        let mut data_bytes = vec![];
+        for elem in elements.iter() {
+            let v = objint::to_int(vm, elem, 10)?;
+            data_bytes.push(v.to_u8().unwrap());
+        }
+        data_bytes
+    // return Err(vm.new_type_error("Cannot construct bytes".to_string()));
+    } else {
+        vec![]
+    };
+
+    Ok(PyObject::new(
+        PyObjectKind::Bytes { value: value },
+        cls.clone(),
+    ))
+}
+
+fn bytearray_eq(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+    arg_check!(
+        vm,
+        args,
+        required = [(a, Some(vm.ctx.bytearray_type())), (b, None)]
+    );
+
+    let result = if objtype::isinstance(b, &vm.ctx.bytearray_type()) {
+        get_value(a).to_vec() == get_value(b).to_vec()
+    } else {
+        false
+    };
+    Ok(vm.ctx.new_bool(result))
+}
+
+/*
+fn set_value(obj: &PyObjectRef, value: Vec<u8>) {
+    obj.borrow_mut().kind = PyObjectKind::Bytes { value };
+}
+*/
+
+fn bytearray_repr(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+    arg_check!(vm, args, required = [(obj, Some(vm.ctx.bytearray_type()))]);
+    let data = get_value(obj);
+    let data: Vec<String> = data.iter().map(|b| format!("\\x{:02x}", b)).collect();
+    let data = data.join("");
+    Ok(vm.new_str(format!("bytearray(b'{}')", data)))
+}
diff --git a/vm/src/obj/objbytes.rs b/vm/src/obj/objbytes.rs
@@ -1,40 +1,52 @@
 use super::super::pyobject::{
-    AttributeProtocol, PyContext, PyFuncArgs, PyObjectKind, PyObjectRef, PyResult, TypeProtocol,
+    AttributeProtocol, PyContext, PyFuncArgs, PyObject, PyObjectKind, PyObjectRef, PyResult,
+    TypeProtocol,
 };
 use super::super::vm::VirtualMachine;
 use super::objint;
-use super::objlist;
 use super::objtype;
 use num_traits::ToPrimitive;
+use std::cell::Ref;
+use std::ops::Deref;
 // Binary data support
 
 // Fill bytes class methods:
 pub fn init(context: &PyContext) {
     let ref bytes_type = context.bytes_type;
     bytes_type.set_attr("__eq__", context.new_rustfunc(bytes_eq));
-    bytes_type.set_attr("__init__", context.new_rustfunc(bytes_init));
+    bytes_type.set_attr("__new__", context.new_rustfunc(bytes_new));
     bytes_type.set_attr("__repr__", context.new_rustfunc(bytes_repr));
 }
 
-// __init__ (store value into objectkind)
-fn bytes_init(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+fn bytes_new(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     arg_check!(
         vm,
         args,
-        required = [(zelf, Some(vm.ctx.bytes_type())), (arg, None)]
+        required = [(cls, None)],
+        optional = [(val_option, None)]
     );
-    let val = if objtype::isinstance(arg, &vm.ctx.list_type()) {
+    if !objtype::issubclass(cls, &vm.ctx.bytes_type()) {
+        return Err(vm.new_type_error(format!("{:?} is not a subtype of bytes", cls)));
+    }
+
+    // Create bytes data:
+    let value = if let Some(ival) = val_option {
+        let elements = vm.extract_elements(ival)?;
         let mut data_bytes = vec![];
-        for elem in objlist::get_elements(arg).iter() {
+        for elem in elements.iter() {
             let v = objint::to_int(vm, elem, 10)?;
             data_bytes.push(v.to_u8().unwrap());
         }
         data_bytes
+    // return Err(vm.new_type_error("Cannot construct bytes".to_string()));
     } else {
-        return Err(vm.new_type_error("Cannot construct bytes".to_string()));
+        vec![]
     };
-    set_value(zelf, val);
-    Ok(vm.get_none())
+
+    Ok(PyObject::new(
+        PyObjectKind::Bytes { value: value },
+        cls.clone(),
+    ))
 }
 
 fn bytes_eq(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
@@ -45,29 +57,27 @@ fn bytes_eq(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     );
 
     let result = if objtype::isinstance(b, &vm.ctx.bytes_type()) {
-        get_value(a) == get_value(b)
+        get_value(a).to_vec() == get_value(b).to_vec()
     } else {
         false
     };
     Ok(vm.ctx.new_bool(result))
 }
 
-pub fn get_value(obj: &PyObjectRef) -> Vec<u8> {
-    if let PyObjectKind::Bytes { value } = &obj.borrow().kind {
-        value.clone()
-    } else {
-        panic!("Inner error getting int {:?}", obj);
-    }
-}
-
-fn set_value(obj: &PyObjectRef, value: Vec<u8>) {
-    obj.borrow_mut().kind = PyObjectKind::Bytes { value };
+pub fn get_value<'a>(obj: &'a PyObjectRef) -> impl Deref<Target = Vec<u8>> + 'a {
+    Ref::map(obj.borrow(), |py_obj| {
+        if let PyObjectKind::Bytes { ref value } = py_obj.kind {
+            value
+        } else {
+            panic!("Inner error getting int {:?}", obj);
+        }
+    })
 }
 
 fn bytes_repr(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     arg_check!(vm, args, required = [(obj, Some(vm.ctx.bytes_type()))]);
     let data = get_value(obj);
-    let data: Vec<String> = data.into_iter().map(|b| format!("\\x{:02x}", b)).collect();
+    let data: Vec<String> = data.iter().map(|b| format!("\\x{:02x}", b)).collect();
     let data = data.join("");
     Ok(vm.new_str(format!("b'{}'", data)))
 }
diff --git a/vm/src/obj/objlist.rs b/vm/src/obj/objlist.rs
@@ -68,12 +68,10 @@ fn list_new(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
         return Err(vm.new_type_error(format!("{:?} is not a subtype of list", cls)));
     }
 
-    let elements = match iterable {
-        None => vec![],
-        Some(iterable) => {
-            let iterator = objiter::get_iter(vm, iterable)?;
-            objiter::get_all(vm, &iterator)?
-        }
+    let elements = if let Some(iterable) = iterable {
+        vm.extract_elements(iterable)?
+    } else {
+        vec![]
     };
 
     Ok(PyObject::new(
diff --git a/vm/src/obj/objtuple.rs b/vm/src/obj/objtuple.rs
@@ -1,5 +1,6 @@
 use super::super::pyobject::{
-    AttributeProtocol, PyContext, PyFuncArgs, PyObjectKind, PyObjectRef, PyResult, TypeProtocol,
+    AttributeProtocol, PyContext, PyFuncArgs, PyObject, PyObjectKind, PyObjectRef, PyResult,
+    TypeProtocol,
 };
 use super::super::vm::VirtualMachine;
 use super::objbool;
@@ -55,6 +56,30 @@ fn tuple_len(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     Ok(vm.context().new_int(elements.len().to_bigint().unwrap()))
 }
 
+fn tuple_new(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+    arg_check!(
+        vm,
+        args,
+        required = [(cls, None)],
+        optional = [(iterable, None)]
+    );
+
+    if !objtype::issubclass(cls, &vm.ctx.tuple_type()) {
+        return Err(vm.new_type_error(format!("{:?} is not a subtype of tuple", cls)));
+    }
+
+    let elements = if let Some(iterable) = iterable {
+        vm.extract_elements(iterable)?
+    } else {
+        vec![]
+    };
+
+    Ok(PyObject::new(
+        PyObjectKind::Tuple { elements: elements },
+        cls.clone(),
+    ))
+}
+
 fn tuple_repr(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     arg_check!(vm, args, required = [(zelf, Some(vm.ctx.tuple_type()))]);
 
@@ -120,5 +145,6 @@ pub fn init(context: &PyContext) {
     tuple_type.set_attr("__getitem__", context.new_rustfunc(tuple_getitem));
     tuple_type.set_attr("__hash__", context.new_rustfunc(tuple_hash));
     tuple_type.set_attr("__len__", context.new_rustfunc(tuple_len));
+    tuple_type.set_attr("__new__", context.new_rustfunc(tuple_new));
     tuple_type.set_attr("__repr__", context.new_rustfunc(tuple_repr));
 }
diff --git a/vm/src/pyobject.rs b/vm/src/pyobject.rs
@@ -2,6 +2,7 @@ use super::bytecode;
 use super::exceptions;
 use super::frame::Frame;
 use super::obj::objbool;
+use super::obj::objbytearray;
 use super::obj::objbytes;
 use super::obj::objcomplex;
 use super::obj::objdict;
@@ -63,6 +64,7 @@ pub struct PyContext {
     pub float_type: PyObjectRef,
     pub complex_type: PyObjectRef,
     pub bytes_type: PyObjectRef,
+    pub bytearray_type: PyObjectRef,
     pub bool_type: PyObjectRef,
     pub true_value: PyObjectRef,
     pub false_value: PyObjectRef,
@@ -137,6 +139,7 @@ impl PyContext {
         let float_type = create_type("float", &type_type, &object_type, &dict_type);
         let complex_type = create_type("complex", &type_type, &object_type, &dict_type);
         let bytes_type = create_type("bytes", &type_type, &object_type, &dict_type);
+        let bytearray_type = create_type("bytearray", &type_type, &object_type, &dict_type);
         let tuple_type = create_type("tuple", &type_type, &object_type, &dict_type);
         let iter_type = create_type("iter", &type_type, &object_type, &dict_type);
         let bool_type = create_type("bool", &type_type, &int_type, &dict_type);
@@ -162,6 +165,7 @@ impl PyContext {
             float_type: float_type,
             complex_type: complex_type,
             bytes_type: bytes_type,
+            bytearray_type: bytearray_type,
             list_type: list_type,
             set_type: set_type,
             bool_type: bool_type,
@@ -193,6 +197,7 @@ impl PyContext {
         objfloat::init(&context);
         objcomplex::init(&context);
         objbytes::init(&context);
+        objbytearray::init(&context);
         objstr::init(&context);
         objtuple::init(&context);
         objiter::init(&context);
@@ -217,6 +222,10 @@ impl PyContext {
         self.bytes_type.clone()
     }
 
+    pub fn bytearray_type(&self) -> PyObjectRef {
+        self.bytearray_type.clone()
+    }
+
     pub fn list_type(&self) -> PyObjectRef {
         self.list_type.clone()
     }
diff --git a/vm/src/stdlib/pystruct.rs b/vm/src/stdlib/pystruct.rs
@@ -313,7 +313,7 @@ fn struct_unpack(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     let fmt_str = objstr::get_value(&fmt);
 
     let codes = parse_format_string(fmt_str);
-    let data = objbytes::get_value(buffer);
+    let data = objbytes::get_value(buffer).to_vec();
     let mut rdr = Cursor::new(data);
 
     let mut items = vec![];
