py/objint: In to_bytes(), allow length arg to be any int and check sign.

pauljay · Jun 15, 2017 · e269cab · e269cab
1 parent 8c5632a
commit e269cab
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 1 deletion.
diff --git a/py/objint.c b/py/objint.c
@@ -435,7 +435,10 @@ STATIC mp_obj_t int_to_bytes(size_t n_args, const mp_obj_t *args) {
     // TODO: Support signed param (assumes signed=False)
     (void)n_args;
 
-    mp_uint_t len = MP_OBJ_SMALL_INT_VALUE(args[1]);
+    mp_int_t len = mp_obj_get_int(args[1]);
+    if (len < 0) {
+        mp_raise_ValueError(NULL);
+    }
     bool big_endian = args[2] != MP_OBJ_NEW_QSTR(MP_QSTR_little);
 
     vstr_t vstr;

diff --git a/tests/basics/int_bytes.py b/tests/basics/int_bytes.py
@@ -14,3 +14,9 @@
 print((100).to_bytes(10, "big"))
 print(int.from_bytes(b"\0\0\0\0\0\0\0\0\0\x01", "big"))
 print(int.from_bytes(b"\x01\0", "big"))
+
+# negative number of bytes should raise an error
+try:
+    (1).to_bytes(-1, "little")
+except ValueError:
+    print("ValueError")