Add methods for serialization

Author: UserUnknownFactor

kaitaistruct.py

@@ -125,6 +125,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 +141,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 +163,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 +195,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 +217,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 +245,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 +261,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
     # ========================================================================
@@ -293,6 +347,21 @@ def read_bytes(self, n):
             )
         return r
 
+    def write_bytes(self, data, pad=0, padding=b'\0'):
+        if data is None:
+            return
+        nb = len(data)
+        if nb == 0 and pad < 1:
+            return
+        if self._io.write(data) != nb:
+            raise Exception(
+                "not all of %d bytes written" %
+                (nb,)
+            )
+        if pad > 0 and padding is not None:
+            self._io.write(padding * pad)
+        return
+
     def read_bytes_full(self):
         return self._io.read()
 
@@ -317,6 +386,9 @@ def read_bytes_term(self, term, include_term, consume_term, eos_error):
             else:
                 r += c
 
+    def write_bytes_term(self, data, term=b'\0'):
+        self.write_bytes(data, pad=1, padding=term)
+
     def ensure_fixed_contents(self, expected):
         actual = self._io.read(len(expected))
         if actual != expected:
@@ -327,7 +399,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