From 2633f9f844c05180a5ce3c0cc3ca077166904e04 Mon Sep 17 00:00:00 2001
From: Aviv Palivoda <palaviv@gmail.com>
Date: Thu, 6 Feb 2020 21:20:39 +0200
Subject: [PATCH] Add struct.calcsize

---
 tests/snippets/stdlib_struct.py |  3 +++
 vm/src/stdlib/pystruct.rs       | 21 +++++++++++++++++++++
 2 files changed, 24 insertions(+)

diff --git a/tests/snippets/stdlib_struct.py b/tests/snippets/stdlib_struct.py
index e7be196c9..100c60bcf 100644
--- a/tests/snippets/stdlib_struct.py
+++ b/tests/snippets/stdlib_struct.py
@@ -41,3 +41,6 @@ data = struct.pack('B1B', 65, 66)
 
 with assert_raises(Exception):
   struct.pack('<IH', "14", 12)
+
+assert struct.calcsize("B") == 1
+assert struct.calcsize("<L4B") == 8
diff --git a/vm/src/stdlib/pystruct.rs b/vm/src/stdlib/pystruct.rs
index 34baea298..2c61784da 100644
--- a/vm/src/stdlib/pystruct.rs
+++ b/vm/src/stdlib/pystruct.rs
@@ -42,6 +42,20 @@ struct FormatCode {
     code: char,
 }
 
+impl FormatCode {
+    fn size(&self) -> usize {
+        match self.code {
+            'b' | 'B' | '?' => 1,
+            'h' | 'H' => 2,
+            'i' | 'l' | 'I' | 'L' | 'f' => 4,
+            'q' | 'Q' | 'd' => 8,
+            c => {
+                panic!("Unsupported format code {:?}", c);
+            }
+        }
+    }
+}
+
 fn parse_format_string(fmt: String) -> Result<FormatSpec, String> {
     let mut chars = fmt.chars().peekable();
 
@@ -402,6 +416,12 @@ where
     }
 }
 
+fn struct_calcsize(fmt: PyStringRef, vm: &VirtualMachine) -> PyResult<usize> {
+    let fmt_str = fmt.as_str().to_owned();
+    let format_spec = parse_format_string(fmt_str).map_err(|e| vm.new_value_error(e))?;
+    Ok(format_spec.codes.iter().map(|code| code.size()).sum())
+}
+
 pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
     let ctx = &vm.ctx;
 
@@ -410,6 +430,7 @@ pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
     py_module!(vm, "struct", {
         "pack" => ctx.new_function(struct_pack),
         "unpack" => ctx.new_function(struct_unpack),
+        "calcsize" => ctx.new_function(struct_calcsize),
         "error" => struct_error,
     })
 }