Improved file reading: started using memory-mapping files and added from_any method

Author: KOLANICH

kaitaistruct.py

@@ -2,6 +2,9 @@
 import sys
 import struct
 from io import open, BytesIO, SEEK_CUR, SEEK_END  # noqa
+from _io import _IOBase
+import mmap
+from pathlib import Path
 
 PY2 = sys.version_info[0] == 2
 
@@ -14,8 +17,40 @@
 #
 __version__ = '0.9'
 
+class _KaitaiFromCtor(object):
+    """Adds to class methods to construct it from ctor"""
+    @classmethod
+    def from_file(cls, file, map=True):
+        if isinstance(file, str):
+            return cls.from_file(Path(file), map)
+        elif isinstance(file, Path):
+            with file.open("rb") as f:
+                return cls.from_file(f, map)
+        elif isinstance(file, _IOBase):
+            if map:
+                with mmap.mmap(file.fileno(), 0, access=mmap.ACCESS_READ) as buf:
+                    return cls.from_io(buf)
+            else:
+                return cls.from_io(file)
+        else:
+            raise Exception("file argument must be either file path (either str or Path) or a file object")
+
+    @classmethod
+    def from_any(cls, data, map=True):
+        if isinstance(data, KaitaiStream):
+            return cls(data)
+        elif isinstance(data, (str, Path, _IOBase)):
+            return cls.from_file(data, map)
+        elif isinstance(data, (bytes, bytearray)):
+            return cls.from_bytes(data)
+        else:
+            return cls.from_io(data)
+
+    @classmethod
+    def from_bytes(cls, buf):
+        return cls.from_io(BytesIO(buf))
 
-class KaitaiStruct(object):
+class KaitaiStruct(_KaitaiFromCtor):
     def __init__(self, stream):
         self._io = stream
 
@@ -28,26 +63,12 @@ def __exit__(self, *args, **kwargs):
     def close(self):
         self._io.close()
 
-    @classmethod
-    def from_file(cls, filename):
-        f = open(filename, 'rb')
-        try:
-            return cls(KaitaiStream(f))
-        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 from_io(cls, io):
         return cls(KaitaiStream(io))
 
 
-class KaitaiStream(object):
+class KaitaiStream(_KaitaiFromCtor):
     def __init__(self, io):
         self._io = io
         self.align_to_byte()
@@ -61,6 +82,10 @@ def __exit__(self, *args, **kwargs):
     def close(self):
         self._io.close()
 
+    @classmethod
+    def from_io(cls, io):
+        return cls(io)
+
     # ========================================================================
     # Stream positioning
     # ========================================================================