modstruct: Refactor to support both LE and BE packed structs.

pfalcon · pfalcon · commit 6582d64d01c4 · 2014-04-11T03:58:49.000+03:00
diff --git a/py/binary.c b/py/binary.c
@@ -77,50 +77,48 @@ mp_obj_t mp_binary_get_val(char typecode, void *p, int index) {
     return MP_OBJ_NEW_SMALL_INT(val);
 }
 
-mp_obj_t mp_binary_get_val_unaligned_le(char typecode, byte **ptr) {
-    machine_int_t val = 0;
+#define is_signed(typecode) (typecode > 'Z')
+mp_obj_t mp_binary_get_val_unaligned(char typecode, byte **ptr) {
+    char type = '<';
     byte *p = *ptr;
-    switch (typecode) {
-        case 'b':
-            val = (int8_t)*p++;
-            break;
-        case BYTEARRAY_TYPECODE:
-        case 'B':
-            val = *p++;
-            break;
-        case 'h':
-            val = (int16_t)((p[1] << 8) | p[0]);
-            break;
-        case 'H':
-            val = (p[1] << 8) | p[0];
+    uint size = 0, align = 0;
+    switch (type) {
+        case '<': case '>':
+            switch (typecode) {
+                case 'b': case 'B':
+                    size = 1; break;
+                case 'h': case 'H':
+                    size = 2; break;
+                case 'i': case 'I':
+                    size = 4; break;
+            }
             break;
-        case 'i':
-        case 'l':
-            val = (p[3] << 24) | (p[2] << 16) | (p[1] << 8) | p[0];
-            *ptr = p + 4;
-            return mp_obj_new_int(val);
-        case 'I':
-        case 'L':
-            val = (p[3] << 24) | (p[2] << 16) | (p[1] << 8) | p[0];
-            *ptr = p + 4;
-            return mp_obj_new_int_from_uint(val);
-#if 0 //TODO
-#if MICROPY_LONGINT_IMPL != MICROPY_LONGINT_IMPL_NONE
-        case 'q':
-        case 'Q':
-            // TODO: Explode API more to cover signedness
-            return mp_obj_new_int_from_ll(((long long*)p)[index]);
-#endif
-#if MICROPY_ENABLE_FLOAT
-        case 'f':
-            return mp_obj_new_float(((float*)p)[index]);
-        case 'd':
-            return mp_obj_new_float(((double*)p)[index]);
-#endif
-#endif
     }
-    *ptr = p;
-    return MP_OBJ_NEW_SMALL_INT(val);
+
+    int delta;
+    if (type == '<') {
+        delta = -1;
+        p += size - 1;
+    } else {
+        delta = 1;
+    }
+
+    machine_int_t val = 0;
+    if (is_signed(typecode) && *p & 0x80) {
+        val = -1;
+    }
+    for (uint i = 0; i < size; i++) {
+        val <<= 8;
+        val |= *p;
+        p += delta;
+    }
+
+    *ptr += size + align;
+    if (is_signed(typecode)) {
+        return mp_obj_new_int(val);
+    } else {
+        return mp_obj_new_int_from_uint(val);
+    }
 }
 
 void mp_binary_set_val(char typecode, void *p, int index, mp_obj_t val_in) {
diff --git a/py/binary.h b/py/binary.h
@@ -4,5 +4,5 @@
 
 int mp_binary_get_size(char typecode);
 mp_obj_t mp_binary_get_val(char typecode, void *p, int index);
-mp_obj_t mp_binary_get_val_unaligned_le(char typecode, byte **ptr);
+mp_obj_t mp_binary_get_val_unaligned(char typecode, byte **ptr);
 void mp_binary_set_val(char typecode, void *p, int index, mp_obj_t val_in);
diff --git a/py/modstruct.c b/py/modstruct.c
@@ -37,7 +37,7 @@ STATIC uint calcsize_items(const char *fmt) {
 STATIC mp_obj_t struct_calcsize(mp_obj_t fmt_in) {
     const char *fmt = mp_obj_str_get_str(fmt_in);
     char fmt_type = get_fmt_type(&fmt);
-    assert(fmt_type == '<'); (void)fmt_type;
+    assert(fmt_type == '<' || fmt_type == '>'); (void)fmt_type;
     machine_uint_t size;
     for (size = 0; *fmt; fmt++) {
         int sz = mp_binary_get_size(*fmt);
@@ -53,15 +53,15 @@ STATIC mp_obj_t struct_unpack(mp_obj_t fmt_in, mp_obj_t data_in) {
     // TODO: "The buffer must contain exactly the amount of data required by the format (len(bytes) must equal calcsize(fmt))."
     const char *fmt = mp_obj_str_get_str(fmt_in);
     char fmt_type = get_fmt_type(&fmt);
-    assert(fmt_type == '<'); (void)fmt_type;
+    assert(fmt_type == '<' || fmt_type == '>'); (void)fmt_type;
     uint size = calcsize_items(fmt);
     mp_obj_tuple_t *res = mp_obj_new_tuple(size, NULL);
     buffer_info_t bufinfo;
     mp_get_buffer_raise(data_in, &bufinfo);
     byte *p = bufinfo.buf;
 
     for (uint i = 0; i < size; i++) {
-        mp_obj_t item = mp_binary_get_val_unaligned_le(*fmt++, &p);
+        mp_obj_t item = mp_binary_get_val_unaligned(*fmt++, &p);
         res->items[i] = item;
     }
     return res;
diff --git a/tests/basics/struct1.py b/tests/basics/struct1.py
@@ -1,3 +1,5 @@
 import struct
 print(struct.calcsize("<bI"))
 print(struct.unpack("<bI", b"\x80\0\0\x01\0"))
+print(struct.calcsize(">bI"))
+print(struct.unpack(">bI", b"\x80\0\0\x01\0"))
