Skip to content

Commit f75052a

Browse files
committed
Complete overhaul of interaction between KaitaiStream and KaitaiStruct
1 parent c6f581a commit f75052a

File tree

1 file changed

+175
-25
lines changed

1 file changed

+175
-25
lines changed

kaitaistruct.py

Lines changed: 175 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,12 @@
1+
import typing
12
import itertools
23
import sys
34
import struct
45
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
510

611
PY2 = sys.version_info[0] == 2
712

@@ -15,51 +20,196 @@
1520
__version__ = '0.9'
1621

1722

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
2155

2256
def __enter__(self):
57+
self._shouldExit = not stream.is_entered
58+
if self._shouldExit:
59+
self._io.__enter__()
2360
return self
2461

2562
def __exit__(self, *args, **kwargs):
26-
self.close()
63+
if self.shouldExit:
64+
self._io.__exit__(*args, **kwargs)
2765

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
3072

3173
@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)
4076

4177
@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)
4480

4581
@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)
4884

4985

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)
54115

55116
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__()
56149
return self
57150

58151
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()
60198

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)
63213

64214
# ========================================================================
65215
# Stream positioning

0 commit comments

Comments
 (0)