diff --git a/kaitaistruct.py b/kaitaistruct.py index c3ce094..d0e7a01 100644 --- a/kaitaistruct.py +++ b/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