Initial struct module

windelbouwman · windelbouwman · commit 5e27cc45b083 · 2018-10-13T22:04:49.000+02:00
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/tests/snippets/basic_types.py b/tests/snippets/basic_types.py
@@ -28,6 +28,9 @@
 
 a = bytes([1, 2, 3])
 print(a)
+b = bytes([1, 2, 3])
+assert a == b
+
 try:
     bytes([object()])
 except TypeError:
diff --git a/tests/snippets/test_struct.py b/tests/snippets/test_struct.py
@@ -0,0 +1,6 @@
+
+import struct
+
+data = struct.pack('IH', 14, 12)
+assert data == bytes([14,0,0,0,12,0])
+
diff --git a/vm/Cargo.toml b/vm/Cargo.toml
@@ -10,3 +10,4 @@ rustpython_parser = {path = "../parser"}
 serde = "1.0.66"
 serde_derive = "1.0.66"
 serde_json = "1.0.26"
+byteorder = "1.2.6"
diff --git a/vm/src/obj/objbytes.rs b/vm/src/obj/objbytes.rs
@@ -10,6 +10,7 @@ use super::objtype;
 // 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("__repr__", context.new_rustfunc(bytes_repr));
 }
@@ -38,6 +39,21 @@ fn bytes_init(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     Ok(vm.get_none())
 }
 
+fn bytes_eq(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+    arg_check!(
+        vm,
+        args,
+        required = [(a, Some(vm.ctx.bytes_type())), (b, None)]
+    );
+
+    let result = if objtype::isinstance(b, vm.ctx.bytes_type()) {
+        get_value(a) == get_value(b)
+    } 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()
diff --git a/vm/src/pyobject.rs b/vm/src/pyobject.rs
@@ -246,6 +246,10 @@ impl PyContext {
         PyObject::new(PyObjectKind::String { value: s }, self.str_type())
     }
 
+    pub fn new_bytes(&self, data: Vec<u8>) -> PyObjectRef {
+        PyObject::new(PyObjectKind::Bytes { value: data }, self.bytes_type())
+    }
+
     pub fn new_bool(&self, b: bool) -> PyObjectRef {
         if b {
             self.true_value.clone()
diff --git a/vm/src/stdlib/mod.rs b/vm/src/stdlib/mod.rs
@@ -3,6 +3,7 @@ mod io;
 mod json;
 mod keyword;
 mod math;
+mod pystruct;
 use std::collections::HashMap;
 
 use super::pyobject::{PyContext, PyObjectRef};
@@ -16,5 +17,6 @@ pub fn get_module_inits() -> HashMap<String, StdlibInitFunc> {
     modules.insert("ast".to_string(), ast::mk_module as StdlibInitFunc);
     modules.insert("keyword".to_string(), keyword::mk_module as StdlibInitFunc);
     modules.insert("math".to_string(), math::mk_module as StdlibInitFunc);
+    modules.insert("struct".to_string(), pystruct::mk_module as StdlibInitFunc);
     modules
 }
diff --git a/vm/src/stdlib/pystruct.rs b/vm/src/stdlib/pystruct.rs
@@ -0,0 +1,118 @@
+/*
+ * Python struct module.
+ *
+ * renamed to pystruct since struct is a rust keyword.
+ *
+ * Use this rust module to do byte packing:
+ * https://docs.rs/byteorder/1.2.6/byteorder/
+ */
+
+extern crate byteorder;
+use self::byteorder::{LittleEndian, WriteBytesExt};
+
+use super::super::obj::{objint, objstr, objtype};
+use super::super::pyobject::{DictProtocol, PyContext, PyFuncArgs, PyObjectRef, PyResult};
+use super::super::VirtualMachine;
+
+#[derive(Debug)]
+struct FormatCode {
+    code: char,
+    size: i32,
+    repeat: i32,
+}
+
+// fn bp_float(v: PyObjectRef) {}
+
+// fn bp_uint() {}
+
+fn parse_format_string(fmt: String) -> Vec<FormatCode> {
+    // First determine "<", ">","!" or "="
+    // TODO
+
+    // Now, analyze struct string furter:
+    let mut codes = vec![];
+    for c in fmt.chars() {
+        match c {
+            'I' => codes.push(FormatCode {
+                code: c,
+                size: 1,
+                repeat: 1,
+            }),
+            'H' => codes.push(FormatCode {
+                code: c,
+                size: 1,
+                repeat: 1,
+            }),
+            _ => {}
+        }
+    }
+    codes
+}
+
+fn struct_pack(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
+    if args.args.len() < 1 {
+        Err(vm.new_type_error(format!(
+            "Expected at least 1 argument (got: {})",
+            args.args.len()
+        )))
+    } else {
+        let fmt_arg = args.args[0].clone();
+        if objtype::isinstance(&fmt_arg, vm.ctx.str_type()) {
+            let fmt_str = objstr::get_value(&fmt_arg);
+
+            let codes = parse_format_string(fmt_str);
+
+            if codes.len() + 1 == args.args.len() {
+                // Create data vector:
+                let mut data = Vec::<u8>::new();
+                // Loop over all opcodes:
+                for (code, arg) in codes.iter().zip(args.args.iter().skip(1)) {
+                    debug!("code: {:?}", code);
+                    match code.code {
+                        'I' => {
+                            if objtype::isinstance(&arg, vm.ctx.int_type()) {
+                                let v = objint::get_value(arg) as u32;
+                                match data.write_u32::<LittleEndian>(v) {
+                                    Ok(_v) => {}
+                                    Err(err) => panic!("Error: {:?}", err),
+                                }
+                            } else {
+                                return Err(vm.new_type_error(format!("Expected int")));
+                            }
+                        }
+                        'H' => {
+                            if objtype::isinstance(&arg, vm.ctx.int_type()) {
+                                let v = objint::get_value(arg) as u16;
+                                match data.write_u16::<LittleEndian>(v) {
+                                    Ok(_v) => {}
+                                    Err(err) => panic!("Error: {:?}", err),
+                                }
+                            } else {
+                                return Err(vm.new_type_error(format!("Expected int")));
+                            }
+                        }
+                        _ => {
+                            panic!("Unsupported format code");
+                        }
+                    }
+                }
+
+                Ok(vm.ctx.new_bytes(data))
+            } else {
+                Err(vm.new_type_error(format!(
+                    "Expected {} arguments (got: {})",
+                    codes.len() + 1,
+                    args.args.len()
+                )))
+            }
+        } else {
+            Err(vm.new_type_error(format!("First argument must be of str type")))
+        }
+    }
+}
+
+pub fn mk_module(ctx: &PyContext) -> PyObjectRef {
+    let py_mod = ctx.new_module(&"struct".to_string(), ctx.new_scope(None));
+    py_mod.set_item("pack", ctx.new_rustfunc(struct_pack));
+    py_mod
+}
