py/binary: Add support for array('q') and array('Q').

dpgeorge · dpgeorge · commit 81794fcd31e6 · 2015-09-01T16:31:48.000+01:00
diff --git a/py/binary.c b/py/binary.c
@@ -130,12 +130,12 @@ mp_obj_t mp_binary_get_val_array(char typecode, void *p, mp_uint_t index) {
             return mp_obj_new_int(((long*)p)[index]);
         case 'L':
             return mp_obj_new_int_from_uint(((unsigned long*)p)[index]);
-#if MICROPY_LONGINT_IMPL != MICROPY_LONGINT_IMPL_NONE
+        #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
+        case 'Q':
+            return mp_obj_new_int_from_ull(((unsigned long long*)p)[index]);
+        #endif
 #if MICROPY_PY_BUILTINS_FLOAT
         case 'f':
             return mp_obj_new_float(((float*)p)[index]);
@@ -316,6 +316,13 @@ void mp_binary_set_val_array(char typecode, void *p, mp_uint_t index, mp_obj_t v
             ((mp_obj_t*)p)[index] = val_in;
             break;
         default:
+            #if MICROPY_LONGINT_IMPL != MICROPY_LONGINT_IMPL_NONE
+            if ((typecode | 0x20) == 'q' && MP_OBJ_IS_TYPE(val_in, &mp_type_int)) {
+                mp_obj_int_to_bytes_impl(val_in, MP_ENDIANNESS_BIG,
+                    sizeof(long long), (byte*)&((long long*)p)[index]);
+                return;
+            }
+            #endif
             mp_binary_set_val_array_from_int(typecode, p, index, mp_obj_get_int(val_in));
     }
 }
@@ -347,13 +354,13 @@ void mp_binary_set_val_array_from_int(char typecode, void *p, mp_uint_t index, m
         case 'L':
             ((unsigned long*)p)[index] = val;
             break;
-#if MICROPY_LONGINT_IMPL != MICROPY_LONGINT_IMPL_NONE
+        #if MICROPY_LONGINT_IMPL != MICROPY_LONGINT_IMPL_NONE
         case 'q':
-        case 'Q':
-            assert(0);
             ((long long*)p)[index] = val;
+        case 'Q':
+            ((unsigned long long*)p)[index] = val;
             break;
-#endif
+        #endif
 #if MICROPY_PY_BUILTINS_FLOAT
         case 'f':
             ((float*)p)[index] = val;
diff --git a/tests/basics/array_q.py b/tests/basics/array_q.py
@@ -0,0 +1,15 @@
+# test array('q') and array('Q')
+
+from array import array
+
+print(array('q'))
+print(array('Q'))
+
+print(array('q', [0]))
+print(array('Q', [0]))
+
+print(array('q', [-2**63, -1, 0, 1, 2, 2**63-1]))
+print(array('Q', [0, 1, 2, 2**64-1]))
+
+print(bytes(array('q', [-1])))
+print(bytes(array('Q', [2**64-1])))
