Resize pack_string with padding

palaviv · palaviv · commit bbb799f2c738 · 2020-03-14T16:25:02.000+02:00
diff --git a/tests/snippets/stdlib_struct.py b/tests/snippets/stdlib_struct.py
@@ -56,3 +56,12 @@ def __index__(self):
 
 data = struct.pack('B', Indexable(65))
 assert data == bytes([65])
+
+data = struct.pack('5s', b"test1")
+assert data == b"test1"
+
+data = struct.pack('3s', b"test2")
+assert data == b"tes"
+
+data = struct.pack('7s', b"test3")
+assert data == b"test3\0\0"
diff --git a/vm/src/stdlib/pystruct.rs b/vm/src/stdlib/pystruct.rs
@@ -338,8 +338,11 @@ fn pack_string(
     data: &mut dyn Write,
     length: usize,
 ) -> PyResult<()> {
-    let v = PyBytesRef::try_from_object(vm, arg.clone())?;
-    match data.write_all(&v.get_value()[..length]) {
+    let mut v = PyBytesRef::try_from_object(vm, arg.clone())?
+        .get_value()
+        .to_vec();
+    v.resize(length, 0);
+    match data.write_all(&v) {
         Ok(_) => Ok(()),
         Err(e) => Err(new_struct_error(vm, format!("{:?}", e))),
     }
