Complete overhaul of interaction between KaitaiStream and KaitaiStruct

Author: KOLANICH

kaitaistruct.py

@@ -1,7 +1,12 @@
+import typing
 import itertools
 import sys
 import struct
 from io import open, BytesIO, SEEK_CUR, SEEK_END  # noqa
+from _io import _IOBase
+import mmap
+from pathlib import Path
+from abc import ABC, abstractmethod
 
 PY2 = sys.version_info[0] == 2
 
@@ -14,52 +19,146 @@
 #
 __version__ = '0.9'
 
-
-class KaitaiStruct(object):
-    def __init__(self, stream):
+class _KaitaiStruct:
+    __slots__ = ("_io",)
+    def __init__(self, stream: "KaitaiStream"):
         self._io = stream
 
+
+class KaitaiStruct(_KaitaiStruct):
+    __slots__ = ("_shouldExit",)
+    def __init__(self, stream: "KaitaiStream"):
+        super.__init__(stream)
+        self._shouldExit = False
+
     def __enter__(self):
+        self._shouldExit = not stream.is_entered
+        if self._shouldExit:
+            self._io.__enter__()
         return self
 
     def __exit__(self, *args, **kwargs):
-        self.close()
-
-    def close(self):
-        self._io.close()
+        if self.shouldExit:
+            self._io.__exit__(*args, **kwargs)
 
     @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
+    def from_file(cls, file: typing.Union[Path, str], use_mmap:bool=True):
+        with cls(KaitaiStream(file, use_mmap=use_mmap)) as res:
+            return res
 
-    @classmethod
-    def from_bytes(cls, buf):
-        return cls(KaitaiStream(BytesIO(buf)))
 
-    @classmethod
-    def from_io(cls, io):
-        return cls(KaitaiStream(io))
+class IKaitaiDownStream(ABC):
+    __slots__ =("_io",)
 
+    def __init__(self, _io: typing.Any):
+        self._io = _io
 
-class KaitaiStream(object):
-    def __init__(self, io):
-        self._io = io
-        self.align_to_byte()
+    @abstractmethod
+    @property
+    def is_entered(self):
+        raise NotImplementedError
+
+    @abstractmethod
+    def __enter__(self):
+        raise NotImplementedError()
+
+    def __exit__(self, *args, **kwargs):
+        if self.is_entered():
+            self._io.__exit__(*args, **kwargs)
+            self._io = None
+
+    def __del__(self):
+        self.__exit__(None, None)
+
+
+class KaitaiFileSyscallDownStream(IKaitaiDownStream):
+    __slots__ = ()
+    def __init__(self, path: typing.Union[Path, str]):
+        super().__init__(Path(path))
+
+    @property
+    def is_entered(self):
+        return isinstance(self._io, _IOBase)
+
+    def __enter__(self):
+        self._io = open(self.path).__enter__()
+        return self
+
+
+class KaitaiFileMapDownStream(IKaitaiFileDownStream):
+    __slots__ = ("file",)
+    def __init__(self, path: Path):
+        super().__init__(None)
+        self.file = KaitaiFileSyscallDownStream(path)
+
+    @property
+    def is_entered(self):
+        return isinstance(self._io, mmap.mmap)
 
     def __enter__(self):
+        self.file = self.file.__enter__()
+        self._io = mmap.mmap(self.file.file.fileno(), 0, access=mmap.ACCESS_READ).__enter__()
         return self
 
     def __exit__(self, *args, **kwargs):
-        self.close()
+        super().__exit__(*args, **kwargs)
+        if self.file is not None:
+            self.file.__exit__(*args, **kwargs)
+            self.file = None
 
-    def close(self):
-        self._io.close()
+
+def getFileDownStream(path: Path, *args, use_mmap: bool=True, **kwargs) -> IKaitaiFileDownStream:
+    if use_mmap:
+        cls = KaitaiFileMapDownStream
+    else:
+        cls = KaitaiFileSyscallDownStream
+    return cls(path, *args, **kwargs)
+
+
+class KaitaiBytesDownStream(IKaitaiBytesDownStream):
+    __slots__ = ()
+    def __init__(self, data: bytes):
+        super().__init__(data)
+
+    @property
+    def is_entered(self):
+        return isinstance(self._io, BytesIO)
+
+    def __enter__(self):
+        self._io = BytesIO(self._io).__enter__()
+        return self
+
+
+downstreamMapping = {
+    bytes: KaitaiBytesDownStream,
+    str: getFileDownStream,
+    Path: getFileDownStream,
+}
+
+
+def get_downstream(x: typing.Union[bytes, str, Path], *args, **kwargs) -> IKaitaiDownStream:
+    return downstreamMapping[type(x)](x, *args, **kwargs)
+
+
+class KaitaiStream():
+    def __init__(self, o: typing.Union[bytes, str, Path]):
+        self._downstream = get_downstream(o)
+        self.align_to_byte()
+
+    @property
+    def _io(self):
+        return self.downstream._io
+
+    def __enter__(self):
+        self._downstream.__enter__()
+        return self
+
+    @property
+    def is_entered(self):
+        return self._downstream is not None and self._downstream.is_entered
+
+    def __exit__(self, *args, **kwargs):
+        self._downstream.__exit__(*args, **kwargs)
 
     # ========================================================================
     # Stream positioning