Skip to content

Commit efb6f6a

Browse files
committed
Implement pascal string packing
1 parent ed685c1 commit efb6f6a

2 files changed

Lines changed: 39 additions & 4 deletions

File tree

Lib/test/test_struct.py

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -334,8 +334,6 @@ def assertStructError(func, *args, **kwargs):
334334
assertStructError(struct.pack, format, 0)
335335
assertStructError(struct.unpack, format, b"")
336336

337-
# TODO: RUSTPYTHON
338-
@unittest.expectedFailure
339337
def test_p_code(self):
340338
# Test p ("Pascal string") code.
341339
for code, input, expected, expectedback in [

vm/src/stdlib/pystruct.rs

Lines changed: 39 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
use byteorder::{ReadBytesExt, WriteBytesExt};
1313
use num_bigint::BigInt;
1414
use num_traits::ToPrimitive;
15+
use std::cmp;
1516
use std::io::{Cursor, Read, Write};
1617
use std::iter::Peekable;
1718

@@ -348,6 +349,24 @@ fn pack_string(
348349
}
349350
}
350351

352+
fn pack_pascal(
353+
vm: &VirtualMachine,
354+
arg: &PyObjectRef,
355+
data: &mut dyn Write,
356+
length: usize,
357+
) -> PyResult<()> {
358+
let mut v = PyBytesRef::try_from_object(vm, arg.clone())?
359+
.get_value()
360+
.to_vec();
361+
let string_length = cmp::min(cmp::min(v.len(), 255), length - 1);
362+
data.write_u8(string_length as u8).unwrap();
363+
v.resize(length - 1, 0);
364+
match data.write_all(&v) {
365+
Ok(_) => Ok(()),
366+
Err(e) => Err(new_struct_error(vm, format!("{:?}", e))),
367+
}
368+
}
369+
351370
fn pack_char(vm: &VirtualMachine, arg: &PyObjectRef, data: &mut dyn Write) -> PyResult<()> {
352371
let v = PyBytesRef::try_from_object(vm, arg.clone())?;
353372
if v.len() == 1 {
@@ -385,10 +404,14 @@ where
385404
'N' | 'P' => pack_usize::<Endianness>,
386405
'f' => pack_f32::<Endianness>,
387406
'd' => pack_f64::<Endianness>,
388-
's' | 'p' => {
407+
's' => {
389408
pack_string(vm, &args[0], data, code.repeat as usize)?;
390409
return Ok(1);
391410
}
411+
'p' => {
412+
pack_pascal(vm, &args[0], data, code.repeat as usize)?;
413+
return Ok(1);
414+
}
392415
'x' => {
393416
for _ in 0..code.repeat as usize {
394417
data.write_u8(0).unwrap();
@@ -557,6 +580,16 @@ fn unpack_string(vm: &VirtualMachine, rdr: &mut dyn Read, length: u32) -> PyResu
557580
Ok(vm.ctx.new_bytes(buf))
558581
}
559582

583+
fn unpack_pascal(vm: &VirtualMachine, rdr: &mut dyn Read, length: u32) -> PyResult {
584+
let mut handle = rdr.take(length as u64);
585+
let mut buf: Vec<u8> = Vec::new();
586+
handle.read_to_end(&mut buf).map_err(|_| {
587+
new_struct_error(vm, format!("unpack requires a buffer of {} bytes", length,))
588+
})?;
589+
let string_length = buf[0] as usize;
590+
Ok(vm.ctx.new_bytes(buf[1..=string_length].to_vec()))
591+
}
592+
560593
fn struct_unpack(fmt: PyStringRef, buffer: PyBytesRef, vm: &VirtualMachine) -> PyResult<PyTuple> {
561594
let fmt_str = fmt.as_str();
562595
let format_spec = FormatSpec::parse(fmt_str).map_err(|e| new_struct_error(vm, e))?;
@@ -592,10 +625,14 @@ where
592625
unpack_empty(vm, rdr, code.repeat);
593626
return Ok(());
594627
}
595-
's' | 'p' => {
628+
's' => {
596629
items.push(unpack_string(vm, rdr, code.repeat)?);
597630
return Ok(());
598631
}
632+
'p' => {
633+
items.push(unpack_pascal(vm, rdr, code.repeat)?);
634+
return Ok(());
635+
}
599636
c => {
600637
panic!("Unsupported format code {:?}", c);
601638
}

0 commit comments

Comments
 (0)