|
| 1 | +import typing |
1 | 2 | import itertools
|
2 | 3 | import sys
|
3 | 4 | import struct
|
4 | 5 | from io import open, BytesIO, SEEK_CUR, SEEK_END # noqa
|
| 6 | +from _io import _IOBase |
| 7 | +import mmap |
| 8 | +from pathlib import Path |
| 9 | +from abc import ABC, abstractmethod |
5 | 10 |
|
6 | 11 | PY2 = sys.version_info[0] == 2
|
7 | 12 |
|
|
15 | 20 | __version__ = '0.9'
|
16 | 21 |
|
17 | 22 |
|
18 |
| -class KaitaiStruct(object): |
19 |
| - def __init__(self, stream): |
20 |
| - self._io = stream |
| 23 | +class _KaitaiStruct_: |
| 24 | + __slots__ = ("_io", "_parent", "_root") |
| 25 | + def __init__(self, _io:"KaitaiStream", _parent:typing.Optiohal["_KaitaiStruct"]=None, _root:typing.Optiohal["_KaitaiStruct"]=None): |
| 26 | + self._io = _io |
| 27 | + self._parent = _parent |
| 28 | + self._root = _root if _root else self |
| 29 | + |
| 30 | + |
| 31 | +class KaitaiParser(ABC): |
| 32 | + __slots__ = () # in fact this class not meant to be instantiated |
| 33 | + |
| 34 | + @abstractmethod |
| 35 | + @classmethod |
| 36 | + def parse(cls, struct:_KaitaiStruct_, io:"KaitaiStream"): |
| 37 | + raise NotlmplementedError() |
| 38 | + |
| 39 | + |
| 40 | +class _KaitaiStruct(_KaitaiStruct_): |
| 41 | + __slots__ = () |
| 42 | + _parser = None # :KaitaiParser |
| 43 | + |
| 44 | + def _read(self): |
| 45 | + self.__class__._parser.parse(self, self._io) |
| 46 | + |
| 47 | + |
| 48 | +class KaitaiStruct(_KaitaiStruct): |
| 49 | + __slots__ = ("_shouldExit",) |
| 50 | + def __init__(self, io: typing.Union["KaitaiStream", Path, bytes, str]): |
| 51 | + if not isinstance(io, KaitaiStream): |
| 52 | + io = KaitaiStream(io) |
| 53 | + super.__init__(io) |
| 54 | + self._shouldExit = False |
21 | 55 |
|
22 | 56 | def __enter__(self):
|
| 57 | + self._shouldExit = not stream.is_entered |
| 58 | + if self._shouldExit: |
| 59 | + self._io.__enter__() |
23 | 60 | return self
|
24 | 61 |
|
25 | 62 | def __exit__(self, *args, **kwargs):
|
26 |
| - self.close() |
| 63 | + if self.shouldExit: |
| 64 | + self._io.__exit__(*args, **kwargs) |
27 | 65 |
|
28 |
| - def close(self): |
29 |
| - self._io.close() |
| 66 | + @classmethod |
| 67 | + def from_any(cls, o: typing.Union[Path, str]) -> "KaitaiStruct": |
| 68 | + with KaitaiStream(o) as io: |
| 69 | + s = cls(io) |
| 70 | + s._read() |
| 71 | + return s |
30 | 72 |
|
31 | 73 | @classmethod
|
32 |
| - def from_file(cls, filename): |
33 |
| - f = open(filename, 'rb') |
34 |
| - try: |
35 |
| - return cls(KaitaiStream(f)) |
36 |
| - except Exception: |
37 |
| - # close file descriptor, then reraise the exception |
38 |
| - f.close() |
39 |
| - raise |
| 74 | + def from_file(cls, file: typing.Union[Path, str, _io._BufferedIOBase], use_mmap:bool=True) -> "KaitaiStruct": |
| 75 | + return cls.from_any(file, use_mmap=use_mmap) |
40 | 76 |
|
41 | 77 | @classmethod
|
42 |
| - def from_bytes(cls, buf): |
43 |
| - return cls(KaitaiStream(BytesIO(buf))) |
| 78 | + def from_bytes(cls, data: bytes) -> "KaitaiStruct": |
| 79 | + return cls.from_any(data) |
44 | 80 |
|
45 | 81 | @classmethod
|
46 |
| - def from_io(cls, io): |
47 |
| - return cls(KaitaiStream(io)) |
| 82 | + def from_io(cls, io: _IOBase) -> "KaitaiStruct": |
| 83 | + return cls.from_any(io) |
48 | 84 |
|
49 | 85 |
|
50 |
| -class KaitaiStream(object): |
51 |
| - def __init__(self, io): |
52 |
| - self._io = io |
53 |
| - self.align_to_byte() |
| 86 | +class IKaitaiDownStream(ABC): |
| 87 | + __slots__ =("_io",) |
| 88 | + |
| 89 | + def __init__(self, _io: typing.Any): |
| 90 | + self._io = _io |
| 91 | + |
| 92 | + @abstractmethod |
| 93 | + @property |
| 94 | + def is_entered(self): |
| 95 | + raise NotImplementedError |
| 96 | + |
| 97 | + @abstractmethod |
| 98 | + def __enter__(self): |
| 99 | + raise NotImplementedError() |
| 100 | + |
| 101 | + def __exit__(self, *args, **kwargs): |
| 102 | + if self.is_entered: |
| 103 | + self._io.__exit__(*args, **kwargs) |
| 104 | + self._io = None |
| 105 | + |
| 106 | + |
| 107 | +class KaitaiIODownStream(IKaitaiDownStream): |
| 108 | + __slots__ = () |
| 109 | + def __init__(self, data: typing.Any): |
| 110 | + super().__init__(data) |
| 111 | + |
| 112 | + @property |
| 113 | + def is_entered(self): |
| 114 | + return isinstance(self._io, _IOBase) |
54 | 115 |
|
55 | 116 | def __enter__(self):
|
| 117 | + if not self.is_entered |
| 118 | + self._io = open(self._io).__enter__() |
| 119 | + return self |
| 120 | + |
| 121 | + |
| 122 | +class KaitaiBytesDownStream(KaitaiIODownStream): |
| 123 | + __slots__ = () |
| 124 | + def __init__(self, data: bytes): |
| 125 | + super().__init__(data) |
| 126 | + |
| 127 | + |
| 128 | +class KaitaiFileSyscallDownStream(KaitaiIODownStream): |
| 129 | + __slots__ = () |
| 130 | + def __init__(self, io: typing.Union[Path, str, _IOBase]): |
| 131 | + if isinstance(io, str): |
| 132 | + io = Path(io) |
| 133 | + super().__init__(io) |
| 134 | + |
| 135 | + |
| 136 | +class KaitaiFileMapDownStream(IKaitaiFileDownStream): |
| 137 | + __slots__ = ("file",) |
| 138 | + def __init__(self, io: typing.Union[Path, str, _IOBase]): |
| 139 | + super().__init__(None) |
| 140 | + self.file = KaitaiFileSyscallDownStream(io) |
| 141 | + |
| 142 | + @property |
| 143 | + def is_entered(self): |
| 144 | + return isinstance(self._io, mmap.mmap) |
| 145 | + |
| 146 | + def __enter__(self): |
| 147 | + self.file = self.file.__enter__() |
| 148 | + self._io = mmap.mmap(self.file.file.fileno(), 0, access=mmap.ACCESS_READ).__enter__() |
56 | 149 | return self
|
57 | 150 |
|
58 | 151 | def __exit__(self, *args, **kwargs):
|
59 |
| - self.close() |
| 152 | + super().__exit__(*args, **kwargs) |
| 153 | + if self.file is not None: |
| 154 | + self.file.__exit__(*args, **kwargs) |
| 155 | + self.file = None |
| 156 | + |
| 157 | + |
| 158 | +def get_file_down_stream(path: Path, *args, use_mmap: bool=True, **kwargs) -> IKaitaiDownStream: |
| 159 | + if use_mmap: |
| 160 | + cls = KaitaiFileMapDownStream |
| 161 | + else: |
| 162 | + cls = KaitaiFileSyscallDownStream |
| 163 | + return cls(path, *args, **kwargs) |
| 164 | + |
| 165 | + |
| 166 | +downstreamMapping = { |
| 167 | + bytes: KaitaiBytesDownStream, |
| 168 | + BytesIO: KaitaiBytesDownStream, |
| 169 | + str: get_file_down_stream, |
| 170 | + Path: get_file_down_stream, |
| 171 | + _io._BufferedIOBase: get_file_down_stream, |
| 172 | +} |
| 173 | + |
| 174 | + |
| 175 | +def get_downstream_ctor(x) -> typing.Type[IKaitaiDownStream]: |
| 176 | + ctor = downstreamMapping.get(t, None) |
| 177 | + if ctor: |
| 178 | + return ctor |
| 179 | + types = t.mro() |
| 180 | + for t1 in types[1:]: |
| 181 | + ctor = downstreamMapping.get(t1, None) |
| 182 | + if ctor: |
| 183 | + downstreamMapping[t] = ctor |
| 184 | + return ctor |
| 185 | + raise TypeError("Unsupported type", t, types) |
| 186 | + |
| 187 | + |
| 188 | +def get_downstream(x: typing.Union[bytes, str, Path], *args, **kwargs) -> IKaitaiDownStream: |
| 189 | + return get_downstream_ctor(type(x))(x, *args, **kwargs) |
| 190 | + |
| 191 | + |
| 192 | +class KaitaiStream(): |
| 193 | + def __init__(self, o: typing.Union[bytes, str, Path, IKaitaiDownStream]): |
| 194 | + if not isinstance(o, IKaitaiDownStream): |
| 195 | + o = get_downstream(o) |
| 196 | + self._downstream = o |
| 197 | + self.align_to_byte() |
60 | 198 |
|
61 |
| - def close(self): |
62 |
| - self._io.close() |
| 199 | + @property |
| 200 | + def _io(self): |
| 201 | + return self._downstream._io |
| 202 | + |
| 203 | + def __enter__(self): |
| 204 | + self._downstream.__enter__() |
| 205 | + return self |
| 206 | + |
| 207 | + @property |
| 208 | + def is_entered(self): |
| 209 | + return self._downstream is not None and self._downstream.is_entered |
| 210 | + |
| 211 | + def __exit__(self, *args, **kwargs): |
| 212 | + self._downstream.__exit__(*args, **kwargs) |
63 | 213 |
|
64 | 214 | # ========================================================================
|
65 | 215 | # Stream positioning
|
|
0 commit comments