Add methods for serialization

Author: UserUnknownFactor

kaitaistruct.py

@@ -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,7 +351,10 @@ 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" %
@@ -291,15 +366,31 @@ def read_bytes(self, n):
                 "requested %d bytes, but got only %d bytes" %
                 (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(
@@ -312,11 +403,14 @@ def read_bytes_term(self, term, include_term, consume_term, eos_error):
                 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:
@@ -327,7 +421,7 @@ def ensure_fixed_contents(self, 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:
             # data[...] must yield an integer, to compare with integer pad_byte
@@ -339,18 +433,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:
             # 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]