py/modstruct: Allow to have "0s" in struct format.

1 files changed, 9 insertions, 15 deletions

diff --git a/py/modstruct.c b/py/modstruct.c
index eabc951ae..2016add17 100644
--- a/py/modstruct.c
+++ b/py/modstruct.c
@@ -103,30 +103,24 @@ STATIC mp_obj_t struct_calcsize(mp_obj_t fmt_in) {
     char fmt_type = get_fmt_type(&fmt);
     mp_uint_t size;
     for (size = 0; *fmt; fmt++) {
-        mp_uint_t align = 1;
         mp_uint_t cnt = 1;
         if (unichar_isdigit(*fmt)) {
             cnt = get_fmt_num(&fmt);
         }
 
-        mp_uint_t sz = 0;
         if (*fmt == 's') {
-            sz = cnt;
-            cnt = 1;
-        }
-
-        while (cnt--) {
-            // If we already have size for 's' case, don't set it again
-            if (sz == 0) {
-                sz = (mp_uint_t)mp_binary_get_size(fmt_type, *fmt, &align);
-            }
+            size += cnt;
+        } else {
+            mp_uint_t align;
+            size_t sz = mp_binary_get_size(fmt_type, *fmt, &align);
             if (sz == 0) {
                 nlr_raise(mp_obj_new_exception_msg(&mp_type_ValueError, "unsupported format"));
             }
-            // Apply alignment
-            size = (size + align - 1) & ~(align - 1);
-            size += sz;
-            sz = 0;
+            while (cnt--) {
+                // Apply alignment
+                size = (size + align - 1) & ~(align - 1);
+                size += sz;
+            }
         }
     }
     return MP_OBJ_NEW_SMALL_INT(size);