KaitaiStream code cleanup (#18)

* KaitaiStream code cleanup

* process_xor_* reimplemented for PY2, short code but no benchmarking

* bytes_terminate, bytes_strip_right PY2 cleanup

* process_xor_* uses some dgladkov's code

Author: arekbulski

kaitaistruct.py

@@ -25,15 +25,13 @@ def __enter__(self):
     def __exit__(self, *args, **kwargs):
         self.close()
 
+    def close(self):
+        self._io.close()
+
     @classmethod
     def from_file(cls, filename):
-        f = open(filename, 'rb')
-        try:
+        with open(filename, 'rb') as f:
             return cls(KaitaiStream(f))
-        except Exception:
-            # close file descriptor, then reraise the exception
-            f.close()
-            raise
 
     @classmethod
     def from_bytes(cls, buf):
@@ -43,9 +41,6 @@ def from_bytes(cls, buf):
     def from_io(cls, io):
         return cls(KaitaiStream(io))
 
-    def close(self):
-        self._io.close()
-
 
 class KaitaiStream(object):
     def __init__(self, io):
@@ -85,19 +80,14 @@ def size(self):
         # current file / StringIO size, thus we use the following
         # trick.
         io = self._io
-
         # Remember our current position
         cur_pos = io.tell()
-
         # Seek to the end of the File object
         io.seek(0, 2)
-
         # Remember position, which is equal to the full length
         full_size = io.tell()
-
         # Seek back to the current position
         io.seek(cur_pos)
-
         return full_size
 
     # ========================================================================
@@ -280,32 +270,30 @@ def ensure_fixed_contents(self, expected):
         actual = self._io.read(len(expected))
         if actual != expected:
             raise Exception(
-                "unexpected fixed contents: got %s, was waiting for %s" %
-                (str(actual), str(expected))
+                "unexpected fixed contents: got %r, was waiting for %r" %
+                (actual, expected)
             )
         return actual
 
     @staticmethod
-    def bytes_strip_right(src, pad_byte):
-        new_len = len(src)
+    def bytes_strip_right(data, pad_byte):
+        new_len = len(data)
         if PY2:
-            data = bytearray(src)
-        else:
-            data = src
+            # data[...] must yield an integer, to compare with integer pad_byte
+            data = bytearray(data)
 
         while new_len > 0 and data[new_len - 1] == pad_byte:
             new_len -= 1
 
         return data[:new_len]
 
     @staticmethod
-    def bytes_terminate(src, term, include_term):
+    def bytes_terminate(data, term, include_term):
         new_len = 0
-        max_len = len(src)
+        max_len = len(data)
         if PY2:
-            data = bytearray(src)
-        else:
-            data = src
+            # 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
@@ -322,26 +310,14 @@ def bytes_terminate(src, term, include_term):
     @staticmethod
     def process_xor_one(data, key):
         if PY2:
-            r = bytearray(data)
-            for i in range(len(r)):
-                r[i] ^= key
-            return bytes(r)
+            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:
-            r = bytearray(data)
-            k = bytearray(key)
-            ki = 0
-            kl = len(k)
-            for i in range(len(r)):
-                r[i] ^= k[ki]
-                ki += 1
-                if ki >= kl:
-                    ki = 0
-            return bytes(r)
+            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)))