Skip to content

Commit a914dab

Browse files
committed
Improved file reading: started using memory-mapping files and added from_any method
1 parent c6f581a commit a914dab

File tree

1 file changed

+42
-16
lines changed

1 file changed

+42
-16
lines changed

kaitaistruct.py

Lines changed: 42 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,9 @@
22
import sys
33
import struct
44
from io import open, BytesIO, SEEK_CUR, SEEK_END # noqa
5+
from _io import _IOBase
6+
import mmap
7+
from pathlib import Path
58

69
PY2 = sys.version_info[0] == 2
710

@@ -14,8 +17,41 @@
1417
#
1518
__version__ = '0.9'
1619

20+
class _KaitaiFromCtor(object):
21+
"""Adds to class methods to construct it from ctor"""
22+
@classmethod
23+
def from_file(cls, file, map=True):
24+
if isinstance(file, str):
25+
return cls.from_file(Path(file), map)
26+
elif isinstance(file, Path):
27+
with file.open("rb") as f:
28+
return cls.from_file(f, map)
29+
elif isinstance(file, _IOBase):
30+
if map:
31+
with mmap.mmap(file.fileno(), 0, access=mmap.ACCESS_READ) as buf:
32+
return cls.from_bytes(buf)
33+
else:
34+
return cls.from_io(file)
35+
else:
36+
raise Exception("file argument must be either file path (either str or Path) or a file object")
37+
38+
@classmethod
39+
def from_any(cls, data, map=True):
40+
"""Parses a Kaitai Struct type"""
41+
if isinstance(data, KaitaiStream):
42+
return cls(data)
43+
elif isinstance(data, (str, Path, _IOBase)):
44+
return cls.from_file(data, map)
45+
elif isinstance(data, (bytes, bytearray)):
46+
return cls.from_bytes(data)
47+
else:
48+
return cls.from_io(data)
49+
50+
@classmethod
51+
def from_bytes(cls, buf):
52+
return cls.from_io(BytesIO(buf))
1753

18-
class KaitaiStruct(object):
54+
class KaitaiStruct(_KaitaiFromCtor):
1955
def __init__(self, stream):
2056
self._io = stream
2157

@@ -28,26 +64,12 @@ def __exit__(self, *args, **kwargs):
2864
def close(self):
2965
self._io.close()
3066

31-
@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
40-
41-
@classmethod
42-
def from_bytes(cls, buf):
43-
return cls(KaitaiStream(BytesIO(buf)))
44-
4567
@classmethod
4668
def from_io(cls, io):
4769
return cls(KaitaiStream(io))
4870

4971

50-
class KaitaiStream(object):
72+
class KaitaiStream(_KaitaiFromCtor):
5173
def __init__(self, io):
5274
self._io = io
5375
self.align_to_byte()
@@ -61,6 +83,10 @@ def __exit__(self, *args, **kwargs):
6183
def close(self):
6284
self._io.close()
6385

86+
@classmethod
87+
def from_io(cls, io):
88+
return cls(io)
89+
6490
# ========================================================================
6591
# Stream positioning
6692
# ========================================================================

0 commit comments

Comments
 (0)