Skip to content

Commit 438dcec

Browse files
authored
Merge pull request RustPython#1752 from palaviv/struct-calcsize
Add struct.calcsize
2 parents f873c0e + 2633f9f commit 438dcec

File tree

2 files changed

+24
-0
lines changed

2 files changed

+24
-0
lines changed

tests/snippets/stdlib_struct.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,3 +41,6 @@
4141

4242
with assert_raises(Exception):
4343
struct.pack('<IH', "14", 12)
44+
45+
assert struct.calcsize("B") == 1
46+
assert struct.calcsize("<L4B") == 8

vm/src/stdlib/pystruct.rs

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,20 @@ struct FormatCode {
4242
code: char,
4343
}
4444

45+
impl FormatCode {
46+
fn size(&self) -> usize {
47+
match self.code {
48+
'b' | 'B' | '?' => 1,
49+
'h' | 'H' => 2,
50+
'i' | 'l' | 'I' | 'L' | 'f' => 4,
51+
'q' | 'Q' | 'd' => 8,
52+
c => {
53+
panic!("Unsupported format code {:?}", c);
54+
}
55+
}
56+
}
57+
}
58+
4559
fn parse_format_string(fmt: String) -> Result<FormatSpec, String> {
4660
let mut chars = fmt.chars().peekable();
4761

@@ -402,6 +416,12 @@ where
402416
}
403417
}
404418

419+
fn struct_calcsize(fmt: PyStringRef, vm: &VirtualMachine) -> PyResult<usize> {
420+
let fmt_str = fmt.as_str().to_owned();
421+
let format_spec = parse_format_string(fmt_str).map_err(|e| vm.new_value_error(e))?;
422+
Ok(format_spec.codes.iter().map(|code| code.size()).sum())
423+
}
424+
405425
pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
406426
let ctx = &vm.ctx;
407427

@@ -410,6 +430,7 @@ pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {
410430
py_module!(vm, "struct", {
411431
"pack" => ctx.new_function(struct_pack),
412432
"unpack" => ctx.new_function(struct_unpack),
433+
"calcsize" => ctx.new_function(struct_calcsize),
413434
"error" => struct_error,
414435
})
415436
}

0 commit comments

Comments
 (0)