Add methods for serialization

Author: UserUnknownFactor

kaitaistruct.py

@@ -3,7 +3,7 @@
 import struct
 from io import open, BytesIO, SEEK_CUR, SEEK_END  # noqa
 
-PY2 = sys.version_info[0] == 2
+PY_OLD = sys.version_info[0] == 2 or (sys.version_info[0] == 3 and sys.version_info[1] < 2)
 
 # Kaitai Struct runtime streaming API version, defined as per PEP-0396
 # standard. Used for two purposes:
@@ -38,14 +38,32 @@ def from_file(cls, filename):
             f.close()
             raise
 
+    @classmethod
+    def to_file(cls, filename, data=None):
+        f = open(filename, 'rb')
+        try:
+            return cls(KaitaiStream(f), _mode='w', _data=data)
+        except Exception:
+            # close file descriptor, then reraise the exception
+            f.close()
+            raise
+
     @classmethod
     def from_bytes(cls, buf):
         return cls(KaitaiStream(BytesIO(buf)))
 
+    @classmethod
+    def to_bytes(cls, buf, data=None):
+        return cls(KaitaiStream(BytesIO(buf)), _mode='w', _data=data)
+
     @classmethod
     def from_io(cls, io):
         return cls(KaitaiStream(io))
 
+    @classmethod
+    def to_io(cls, io, data=None):
+        return cls(KaitaiStream(io), _mode='w', _data=data)
+
 
 class KaitaiStream(object):
     def __init__(self, io):
@@ -125,6 +143,9 @@ def size(self):
     def read_s1(self):
         return KaitaiStream.packer_s1.unpack(self.read_bytes(1))[0]
 
+    def write_s1(self, data):
+        return self.write_bytes(KaitaiStream.packer_s1.pack(data))
+
     # ........................................................................
     # Big-endian
     # ........................................................................
@@ -138,6 +159,15 @@ def read_s4be(self):
     def read_s8be(self):
         return KaitaiStream.packer_s8be.unpack(self.read_bytes(8))[0]
 
+    def write_s2be(self, data):
+        return self.write_bytes(KaitaiStream.packer_s2be.pack(data))
+
+    def write_s4be(self, data):
+        return self.write_bytes(KaitaiStream.packer_s4be.pack(data))
+
+    def write_s8be(self, data):
+        return self.write_bytes(KaitaiStream.packer_s8be.pack(data))
+
     # ........................................................................
     # Little-endian
     # ........................................................................
@@ -151,13 +181,25 @@ def read_s4le(self):
     def read_s8le(self):
         return KaitaiStream.packer_s8le.unpack(self.read_bytes(8))[0]
 
+    def write_s2le(self, data):
+        return self.write_bytes(KaitaiStream.packer_s2le.pack(data))
+
+    def write_s4le(self, data):
+        return self.write_bytes(KaitaiStream.packer_s4le.pack(data))
+
+    def write_s8le(self, data):
+        return self.write_bytes(KaitaiStream.packer_s8le.pack(data))
+
     # ------------------------------------------------------------------------
     # Unsigned
     # ------------------------------------------------------------------------
 
     def read_u1(self):
         return KaitaiStream.packer_u1.unpack(self.read_bytes(1))[0]
 
+    def write_u1(self, data):
+        return self.write_bytes(KaitaiStream.packer_u1.pack(data))
+
     # ........................................................................
     # Big-endian
     # ........................................................................
@@ -171,6 +213,15 @@ def read_u4be(self):
     def read_u8be(self):
         return KaitaiStream.packer_u8be.unpack(self.read_bytes(8))[0]
 
+    def write_u2be(self, data):
+        return self.write_bytes(KaitaiStream.packer_u2be.pack(data))
+
+    def write_u4be(self, data):
+        return self.write_bytes(KaitaiStream.packer_u4be.pack(data))
+
+    def write_u8be(self, data):
+        return self.write_bytes(KaitaiStream.packer_u8be.pack(data))
+
     # ........................................................................
     # Little-endian
     # ........................................................................
@@ -184,6 +235,15 @@ def read_u4le(self):
     def read_u8le(self):
         return KaitaiStream.packer_u8le.unpack(self.read_bytes(8))[0]
 
+    def write_u2le(self, data):
+        return self.write_bytes(KaitaiStream.packer_u2le.pack(data))
+
+    def write_u4le(self, data):
+        return self.write_bytes(KaitaiStream.packer_u4le.pack(data))
+
+    def write_u8le(self, data):
+        return self.write_bytes(KaitaiStream.packer_u8le.pack(data))
+
     # ========================================================================
     # Floating point numbers
     # ========================================================================
@@ -203,6 +263,12 @@ def read_f4be(self):
     def read_f8be(self):
         return KaitaiStream.packer_f8be.unpack(self.read_bytes(8))[0]
 
+    def write_f4be(self, data):
+        return self.write_bytes(KaitaiStream.packer_f4be.pack(data))
+
+    def write_f8be(self, data):
+        return self.write_bytes(KaitaiStream.packer_f8be.pack(data))
+
     # ........................................................................
     # Little-endian
     # ........................................................................
@@ -213,6 +279,12 @@ def read_f4le(self):
     def read_f8le(self):
         return KaitaiStream.packer_f8le.unpack(self.read_bytes(8))[0]
 
+    def write_f4le(self, data):
+        return self.write_bytes(KaitaiStream.packer_f4le.pack(data))
+
+    def write_f8le(self, data):
+        return self.write_bytes(KaitaiStream.packer_f8le.pack(data))
+
     # ========================================================================
     # Unaligned bit values
     # ========================================================================
@@ -279,57 +351,67 @@ def read_bits_int_le(self, n):
     # Byte arrays
     # ========================================================================
 
-    def read_bytes(self, n):
+    def alignment(self, a):
+        return (a - self.pos()) % a
+
+    def read_bytes(self, n, align=0):
         if n < 0:
-            raise ValueError(
-                "requested invalid %d amount of bytes" %
-                (n,)
-            )
+            raise ValueError("%d is invalid amount of bytes" % n)
         r = self._io.read(n)
         if len(r) < n:
-            raise EOFError(
-                "requested %d bytes, but got only %d bytes" %
-                (n, len(r))
-            )
+            raise EOFError("got only %d bytes out of %d requested" % (n, len(r)))
+        if align > 1:
+            self._io.seek(self.alignment(align), 1)
         return r
 
+    def write_bytes(self, data, align=0, pad=0, padding=b'\0'):
+        if data is None:
+            return
+        nb = len(data)
+        if nb == 0 and align < 2 and pad < 1:
+            return
+        if self._io.write(data) != nb:
+            raise Exception("not all bytes written")
+        if pad > 0:
+            self._io.write(padding * pad)
+        if align > 1:
+            self._io.write(padding * self.alignment(align))
+        return
+
     def read_bytes_full(self):
         return self._io.read()
 
-    def read_bytes_term(self, term, include_term, consume_term, eos_error):
+    def read_bytes_term(self, term, include_term=False, consume_term=True, eos_error=True, elem_size=1):
         r = b''
         while True:
-            c = self._io.read(1)
+            c = self._io.read(elem_size)
             if c == b'':
                 if eos_error:
-                    raise Exception(
-                        "end of stream reached, but no terminator %d found" %
-                        (term,)
-                    )
+                    raise Exception("end of stream reached, but no terminator (%d) found" % term)
                 else:
                     return r
             elif ord(c) == term:
                 if include_term:
                     r += c
                 if not consume_term:
-                    self._io.seek(-1, SEEK_CUR)
+                    self._io.seek(-elem_size, SEEK_CUR)
                 return r
             else:
                 r += c
 
+    def write_bytes_term(self, data, term=b'\0', align=0):
+        self.write_bytes(data, align=align, pad=1, padding=term)
+
     def ensure_fixed_contents(self, expected):
         actual = self._io.read(len(expected))
         if actual != expected:
-            raise Exception(
-                "unexpected fixed contents: got %r, was waiting for %r" %
-                (actual, expected)
-            )
+            raise Exception("unexpected fixed contents: got %r, was waiting for %r" % (actual, expected))
         return actual
 
     @staticmethod
-    def bytes_strip_right(data, pad_byte):
+    def bytes_strip_right(data, pad_byte=b'\0'):
         new_len = len(data)
-        if PY2:
+        if PY_OLD:
             # data[...] must yield an integer, to compare with integer pad_byte
             data = bytearray(data)
 
@@ -339,18 +421,18 @@ def bytes_strip_right(data, pad_byte):
         return data[:new_len]
 
     @staticmethod
-    def bytes_terminate(data, term, include_term):
+    def bytes_terminate(data, term, include_term=True, elem_size=1):
         new_len = 0
         max_len = len(data)
-        if PY2:
+        if PY_OLD:
             # data[...] must yield an integer, to compare with integer term
             data = bytearray(data)
 
         while new_len < max_len and data[new_len] != term:
-            new_len += 1
+            new_len += elem_size
 
         if include_term and new_len < max_len:
-            new_len += 1
+            new_len += elem_size
 
         return data[:new_len]
 
@@ -360,14 +442,14 @@ def bytes_terminate(data, term, include_term):
 
     @staticmethod
     def process_xor_one(data, key):
-        if PY2:
+        if PY_OLD:
             return bytes(bytearray(v ^ key for v in bytearray(data)))
         else:
             return bytes(v ^ key for v in data)
 
     @staticmethod
     def process_xor_many(data, key):
-        if PY2:
+        if PY_OLD:
             return bytes(bytearray(a ^ b for a, b in zip(bytearray(data), itertools.cycle(bytearray(key)))))
         else:
             return bytes(a ^ b for a, b in zip(data, itertools.cycle(key)))
@@ -393,10 +475,10 @@ def process_rotate_left(data, amount, group_size):
     # ========================================================================
 
     @staticmethod
-    def int_from_byte(v):
-        if PY2:
+    def int_from_byte(v, signed=False):
+        if PY_OLD:
             return ord(v)
-        return v
+        return int.from_bytes(v, signed=signed)
 
     @staticmethod
     def byte_array_index(data, i):