Implemented + tested dna.Struct.field_get

This commit is contained in:
Sybren A. Stüvel
2018-02-22 14:50:24 +01:00
parent 03bf2ac69f
commit e66b1edaf4
6 changed files with 344 additions and 163 deletions
+9 -56
View File
@@ -29,7 +29,7 @@ import pathlib
import tempfile
import typing
from . import exceptions, dna_io, dna
from . import exceptions, dna_io, dna, header
log = logging.getLogger(__name__)
@@ -94,7 +94,7 @@ class BlendFile:
elif magic != BLENDFILE_MAGIC:
raise exceptions.BlendFileError("File is not a blend file", path)
self.header = BlendFileHeader(self.fileobj, self.raw_filepath)
self.header = header.BlendFileHeader(self.fileobj, self.raw_filepath)
self.block_header_struct = self.header.create_block_header_struct()
self.blocks = []
self.code_index = collections.defaultdict(list)
@@ -171,9 +171,10 @@ class BlendFile:
DNACatalog is a catalog of all information in the DNA1 file-block
"""
self.log.debug("building DNA catalog")
shortstruct = self.header.types.USHORT
shortstruct2 = self.header.types.USHORT2
intstruct = self.header.types.UINT
endian = self.header.endian
shortstruct = endian.USHORT
shortstruct2 = endian.USHORT2
intstruct = endian.UINT
assert intstruct.size == 4
data = self.fileobj.read(block.size)
@@ -188,7 +189,7 @@ class BlendFile:
self.log.debug("building #%d names" % names_len)
for _ in range(names_len):
typename = dna_io.read_data0_offset(data, offset)
typename = endian.read_data0_offset(data, offset)
offset = offset + len(typename) + 1
typenames.append(dna.Name(typename))
@@ -198,7 +199,7 @@ class BlendFile:
offset += 4
self.log.debug("building #%d types" % types_len)
for _ in range(types_len):
dna_type_id = dna_io.read_data0_offset(data, offset)
dna_type_id = endian.read_data0_offset(data, offset)
types.append(dna.Struct(dna_type_id))
offset += len(dna_type_id) + 1
@@ -246,54 +247,6 @@ class BlendFile:
return structs, sdna_index_from_id
class BlendFileHeader:
"""
BlendFileHeader represents the first 12 bytes of a blend file.
it contains information about the hardware architecture, which is relevant
to the structure of the rest of the file.
"""
log = log.getChild('BlendFileHeader')
structure = struct.Struct(b'7s1s1s3s')
def __init__(self, fileobj: typing.BinaryIO, path: pathlib.Path):
self.log.debug("reading blend-file-header %s", path)
header = fileobj.read(self.structure.size)
values = self.structure.unpack(header)
self.magic = values[0]
pointer_size_id = values[1]
if pointer_size_id == b'-':
self.pointer_size = 8
elif pointer_size_id == b'_':
self.pointer_size = 4
else:
raise exceptions.BlendFileError('invalid pointer size %r' % pointer_size_id, path)
endian_id = values[2]
if endian_id == b'v':
self.types = dna_io.LittleEndianTypes
self.endian_str = b'<' # indication for struct.Struct()
elif endian_id == b'V':
self.types = dna_io.BigEndianTypes
self.endian_str = b'>' # indication for struct.Struct()
else:
raise exceptions.BlendFileError('invalid endian indicator %r' % endian_id, path)
version_id = values[3]
self.version = int(version_id)
def create_block_header_struct(self) -> struct.Struct:
"""Create a Struct instance for parsing data block headers."""
return struct.Struct(b''.join((
self.endian_str,
b'4sI',
b'I' if self.pointer_size == 4 else b'Q',
b'II',
)))
class BlendFileBlock:
"""
Instance of a struct.
@@ -421,7 +374,7 @@ class BlendFileBlock:
return dna_struct.field_get(
self.file.header, self.file.handle, path,
default=default,
use_nil=use_nil, use_str=use_str,
nil_terminated=use_nil, as_str=use_str,
)
def get_recursive_iter(self, path, path_root=b"",
+85 -51
View File
@@ -1,6 +1,12 @@
import os
import typing
import os
from . import dna_io, header
# Either a simple path b'propname', or a tuple (b'parentprop', b'actualprop', arrayindex)
FieldPath = typing.Union[bytes, typing.Iterable[typing.Union[bytes, int]]]
class Name:
"""dna.Name is a C-type name stored in the DNA as bytes."""
@@ -74,29 +80,52 @@ class Field:
class Struct:
"""dna.Struct is a C-type structure stored in the DNA."""
def __init__(self, dna_type_id: bytes):
def __init__(self, dna_type_id: bytes, size: int = None):
"""
:param dna_type_id: name of the struct in C, like b'AlembicObjectPath'.
:param size: only for unit tests; typically set after construction by
BlendFile.decode_structs(). If not set, it is calculated on the fly
when struct.size is evaluated, based on the available fields.
"""
self.dna_type_id = dna_type_id
self._size = size
self._fields = []
self._fields_by_name = {}
def __repr__(self):
return '%s(%r)' % (type(self).__qualname__, self.dna_type_id)
@property
def size(self) -> int:
if self._size is None:
if not self._fields:
raise ValueError('Unable to determine size of fieldless %r' % self)
last_field = max(self._fields, key=lambda f: f.offset)
self._size = last_field.offset + last_field.size
return self._size
@size.setter
def size(self, new_size: int):
self._size = new_size
def append_field(self, field: Field):
self._fields.append(field)
self._fields_by_name[field.name.name_only] = field
def field_from_path(self,
pointer_size: int,
path: typing.Union[bytes, typing.Iterable[typing.Union[bytes, int]]]) \
-> typing.Tuple[typing.Optional[Field], int]:
path: FieldPath) \
-> typing.Tuple[Field, int]:
"""
Support lookups as bytes or a tuple of bytes and optional index.
C style 'id.name' --> (b'id', b'name')
C style 'array[4]' --> (b'array', 4)
:returns: the field itself, and its offset taking into account the optional index.
:returns: the field itself, and its offset taking into account the
optional index. The offset is relative to the start of the struct,
i.e. relative to the BlendFileBlock containing the data.
:raises KeyError: if the field does not exist.
"""
if isinstance(path, (tuple, list)):
name = path[0]
@@ -124,72 +153,77 @@ class Struct:
return field.dna_type.field_from_path(pointer_size, name_tail)
offset = field.offset
# fileobj.seek(field.offset, os.SEEK_CUR)
if index:
if field.name.is_pointer:
index_offset = pointer_size * index
else:
index_offset = field.dna_type.size * index
if index_offset >= field.size:
raise OverflowError('path %r is out of bounds of its DNA type' % path)
# fileobj.seek(index_offset, os.SEEK_CUR)
raise OverflowError('path %r is out of bounds of its DNA type %s' %
(path, field.dna_type))
offset += index_offset
return field, offset
def field_get(self, header, handle, path,
def field_get(self,
file_header: header.BlendFileHeader,
fileobj: typing.BinaryIO,
path: FieldPath,
default=...,
use_nil=True, use_str=True,
nil_terminated=True,
as_str=True,
):
field = self.field_from_path(header, handle, path)
if field is None:
if default is not ...:
return default
else:
raise KeyError("%r not found in %r (%r)" %
(
path, [f.dna_name.name_only for f in self._fields],
self.dna_type_id))
"""Read the value of the field from the blend file.
Assumes the file pointer of `fileobj` is seek()ed to the start of the
struct on disk (e.g. the start of the BlendFileBlock containing the
data).
"""
try:
field, offset = self.field_from_path(file_header.pointer_size, path)
except KeyError:
if default is ...:
raise
return default
fileobj.seek(offset, os.SEEK_CUR)
dna_type = field.dna_type
dna_name = field.dna_name
dna_size = field.dna_size
dna_name = field.name
types = file_header.endian
# Some special cases (pointers, strings/bytes)
if dna_name.is_pointer:
return DNA_IO.read_pointer(handle, header)
elif dna_type.dna_type_id == b'int':
if dna_name.array_size > 1:
return [DNA_IO.read_int(handle, header) for i in range(dna_name.array_size)]
return DNA_IO.read_int(handle, header)
elif dna_type.dna_type_id == b'short':
if dna_name.array_size > 1:
return [DNA_IO.read_short(handle, header) for i in range(dna_name.array_size)]
return DNA_IO.read_short(handle, header)
elif dna_type.dna_type_id == b'uint64_t':
if dna_name.array_size > 1:
return [DNA_IO.read_ulong(handle, header) for i in range(dna_name.array_size)]
return DNA_IO.read_ulong(handle, header)
elif dna_type.dna_type_id == b'float':
if dna_name.array_size > 1:
return [DNA_IO.read_float(handle, header) for i in range(dna_name.array_size)]
return DNA_IO.read_float(handle, header)
elif dna_type.dna_type_id == b'char':
if dna_size == 1:
return types.read_pointer(fileobj, file_header.pointer_size)
if dna_type.dna_type_id == b'char':
if field.size == 1:
# Single char, assume it's bitflag or int value, and not a string/bytes data...
return DNA_IO.read_char(handle, header)
if use_str:
if use_nil:
return DNA_IO.read_string0(handle, dna_name.array_size)
else:
return DNA_IO.read_string(handle, dna_name.array_size)
return types.read_char(fileobj)
if nil_terminated:
data = types.read_bytes0(fileobj, dna_name.array_size)
else:
if use_nil:
return DNA_IO.read_bytes0(handle, dna_name.array_size)
else:
return DNA_IO.read_bytes(handle, dna_name.array_size)
else:
data = fileobj.read(dna_name.array_size)
if as_str:
return data.decode('utf8')
return data
simple_readers = {
b'int': types.read_int,
b'short': types.read_short,
b'uint64_t': types.read_ulong,
b'float': types.read_float,
}
try:
simple_reader = simple_readers[dna_type.dna_type_id]
except KeyError:
raise NotImplementedError("%r exists but isn't pointer, can't resolve field %r" %
(path, dna_name.name_only), dna_name, dna_type)
if dna_name.array_size > 1:
return [simple_reader(fileobj) for _ in range(dna_name.array_size)]
return simple_reader(fileobj)
def field_set(self, header, handle, path, value):
assert (type(path) == bytes)
+40 -48
View File
@@ -5,7 +5,7 @@ import typing
class EndianIO:
UCHAR = struct.Struct(b'<b')
UCHAR = struct.Struct(b'<B')
USHORT = struct.Struct(b'<H')
USHORT2 = struct.Struct(b'<HH') # two shorts in a row
SSHORT = struct.Struct(b'<h')
@@ -48,17 +48,47 @@ class EndianIO:
return cls._read(fileobj, cls.ULONG)
@classmethod
def read_pointer(cls, fileobj: typing.BinaryIO, fileheader):
"""Read a pointer from a file.
def read_pointer(cls, fileobj: typing.BinaryIO, pointer_size: int):
"""Read a pointer from a file."""
The pointer size is given by the header (BlendFileHeader).
"""
if fileheader.pointer_size == 4:
if pointer_size == 4:
return cls.read_uint(fileobj)
if fileheader.pointer_size == 8:
if pointer_size == 8:
return cls.read_ulong(fileobj)
raise ValueError('unsupported pointer size %d' % fileheader.pointer_size)
raise ValueError('unsupported pointer size %d' % pointer_size)
@classmethod
def write_string(cls, fileobj: typing.BinaryIO, astring: str, fieldlen: int):
assert isinstance(astring, str)
# TODO: truncate the string on a UTF-8 character boundary to avoid creating invalid UTF-8.
cls.write_bytes(fileobj, astring.encode('utf-8'), fieldlen)
@classmethod
def write_bytes(cls, fileobj: typing.BinaryIO, data: bytes, fieldlen: int):
assert isinstance(data, (bytes, bytearray))
if len(data) >= fieldlen:
to_write = data[0:fieldlen]
else:
to_write = data + b'\0'
fileobj.write(to_write)
@classmethod
def read_bytes0(cls, fileobj, length):
data = fileobj.read(length)
return cls.read_data0(data)
@classmethod
def read_data0_offset(cls, data, offset):
add = data.find(b'\0', offset) - offset
return data[offset:offset + add]
@classmethod
def read_data0(cls, data):
add = data.find(b'\0')
if add < 0:
return data
return data[:add]
class LittleEndianTypes(EndianIO):
@@ -66,7 +96,7 @@ class LittleEndianTypes(EndianIO):
class BigEndianTypes(LittleEndianTypes):
UCHAR = struct.Struct(b'>b')
UCHAR = struct.Struct(b'>B')
USHORT = struct.Struct(b'>H')
USHORT2 = struct.Struct(b'>HH') # two shorts in a row
SSHORT = struct.Struct(b'>h')
@@ -74,41 +104,3 @@ class BigEndianTypes(LittleEndianTypes):
SINT = struct.Struct(b'>i')
FLOAT = struct.Struct(b'>f')
ULONG = struct.Struct(b'>Q')
def write_string(fileobj: typing.BinaryIO, astring: str, fieldlen: int):
assert isinstance(astring, str)
write_bytes(fileobj, astring.encode('utf-8'), fieldlen)
def write_bytes(fileobj: typing.BinaryIO, data: bytes, fieldlen: int):
assert isinstance(data, (bytes, bytearray))
if len(data) >= fieldlen:
to_write = data[0:fieldlen]
else:
to_write = data + b'\0'
fileobj.write(to_write)
def read_bytes0(fileobj, length):
data = fileobj.read(length)
return read_data0(data)
def read_string(fileobj, length):
return fileobj.read(length).decode('utf-8')
def read_string0(fileobj, length):
return read_bytes0(fileobj, length).decode('utf-8')
def read_data0_offset(data, offset):
add = data.find(b'\0', offset) - offset
return data[offset:offset + add]
def read_data0(data):
add = data.find(b'\0')
return data[:add]
+55
View File
@@ -0,0 +1,55 @@
import logging
import pathlib
import struct
import typing
from . import dna_io, exceptions
log = logging.getLogger(__name__)
class BlendFileHeader:
"""
BlendFileHeader represents the first 12 bytes of a blend file.
it contains information about the hardware architecture, which is relevant
to the structure of the rest of the file.
"""
structure = struct.Struct(b'7s1s1s3s')
def __init__(self, fileobj: typing.BinaryIO, path: pathlib.Path):
log.debug("reading blend-file-header %s", path)
header = fileobj.read(self.structure.size)
values = self.structure.unpack(header)
self.magic = values[0]
pointer_size_id = values[1]
if pointer_size_id == b'-':
self.pointer_size = 8
elif pointer_size_id == b'_':
self.pointer_size = 4
else:
raise exceptions.BlendFileError('invalid pointer size %r' % pointer_size_id, path)
endian_id = values[2]
if endian_id == b'v':
self.endian = dna_io.LittleEndianTypes
self.endian_str = b'<' # indication for struct.Struct()
elif endian_id == b'V':
self.endian = dna_io.BigEndianTypes
self.endian_str = b'>' # indication for struct.Struct()
else:
raise exceptions.BlendFileError('invalid endian indicator %r' % endian_id, path)
version_id = values[3]
self.version = int(version_id)
def create_block_header_struct(self) -> struct.Struct:
"""Create a Struct instance for parsing data block headers."""
return struct.Struct(b''.join((
self.endian_str,
b'4sI',
b'I' if self.pointer_size == 4 else b'Q',
b'II',
)))