Fixed issue reading array items and renamed some stuff

This commit is contained in:
Sybren A. Stüvel 2018-02-22 16:14:10 +01:00
parent 81dd2c9721
commit 87300df6a3
4 changed files with 60 additions and 19 deletions

View File

@ -100,7 +100,7 @@ class BlendFile:
self.code_index = collections.defaultdict(list)
self.structs = []
self.sdna_index_from_id = {}
self.block_from_offset = {}
self.block_from_addr = {}
self.load_dna1_block()
@ -123,9 +123,8 @@ class BlendFile:
raise exceptions.NoDNA1Block("No DNA1 block in file, not a valid .blend file",
self.filepath)
# cache (could lazy init, incase we never use?)
self.block_from_offset = {block.addr_old: block for block in self.blocks
if block.code != b'ENDB'}
self.block_from_addr = {block.addr_old: block for block in self.blocks
if block.code != b'ENDB'}
def __repr__(self):
clsname = self.__class__.__qualname__
@ -149,7 +148,7 @@ class BlendFile:
# same as looking looping over all blocks,
# then checking ``block.addr_old == offset``
assert (type(offset) is int)
return self.block_from_offset.get(offset)
return self.block_from_addr.get(offset)
def close(self):
"""Close the blend file.
@ -354,16 +353,33 @@ class BlendFileBlock:
return self.bfile.fileobj.tell(), field.dna_name.array_size
def get(self, path,
def get(self,
path: dna.FieldPath,
default=...,
sdna_index_refine=None,
use_nil=True, use_str=True,
null_terminated: typing.Optional[bool]=None,
as_str=True,
base_index=0,
):
"""Read a property and return the value.
:param path: name of the property (like `b'loc'`), tuple of names
to read a sub-property (like `(b'id', b'name')`), or tuple of
name and index to read one item from an array (like
`(b'loc', 2)`)
:param default: The value to return when the field does not exist.
Use Ellipsis (the default value) to raise a KeyError instead.
:param null_terminated: Only used when reading bytes or strings. When
True, stops reading at the first zero byte. Defaults to the same
value as `as_str`, as this is mostly used for string data.
:param as_str: When True, automatically decode bytes to string
(assumes UTF-8 encoding).
"""
ofs = self.file_offset
if base_index != 0:
assert (base_index < self.count)
if base_index >= self.count:
raise OverflowError('%r: index %d overflows size %d' %
(self, base_index, self.count))
ofs += (self.size // self.count) * base_index
self.bfile.fileobj.seek(ofs, os.SEEK_SET)
@ -376,7 +392,7 @@ class BlendFileBlock:
return dna_struct.field_get(
self.bfile.header, self.bfile.fileobj, path,
default=default,
nil_terminated=use_nil, as_str=use_str,
null_terminated=null_terminated, as_str=as_str,
)
def get_recursive_iter(self, path, path_root=b"",
@ -476,7 +492,7 @@ class BlendFileBlock:
# dict like access
def __getitem__(self, item):
return self.get(item, use_str=False)
return self.get(item, as_str=False)
def __setitem__(self, item, value):
self.set(item, value)

View File

@ -130,7 +130,7 @@ class Struct:
i.e. relative to the BlendFileBlock containing the data.
:raises KeyError: if the field does not exist.
"""
if isinstance(path, (tuple, list)):
if isinstance(path, tuple):
name = path[0]
if len(path) >= 2 and not isinstance(path[1], bytes):
name_tail = path[2:]
@ -173,7 +173,7 @@ class Struct:
fileobj: typing.BinaryIO,
path: FieldPath,
default=...,
nil_terminated=True,
null_terminated: typing.Optional[bool]=None,
as_str=True,
):
"""Read the value of the field from the blend file.
@ -181,6 +181,17 @@ class Struct:
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).
:param file_header:
:param fileobj:
:param path:
:param default: The value to return when the field does not exist.
Use Ellipsis (the default value) to raise a KeyError instead.
:param null_terminated: Only used when reading bytes or strings. When
True, stops reading at the first zero byte. Defaults to the same
value as `as_str`, as this is mostly used for string data.
:param as_str: When True, automatically decode bytes to string
(assumes UTF-8 encoding).
"""
try:
field, offset = self.field_from_path(file_header.pointer_size, path)
@ -202,7 +213,7 @@ class Struct:
if field.size == 1:
# Single char, assume it's bitflag or int value, and not a string/bytes data...
return types.read_char(fileobj)
if nil_terminated:
if null_terminated or (null_terminated is None and as_str):
data = types.read_bytes0(fileobj, dna_name.array_size)
else:
data = fileobj.read(dna_name.array_size)
@ -223,6 +234,12 @@ class Struct:
raise NotImplementedError("%r exists but isn't pointer, can't resolve field %r" %
(path, dna_name.name_only), dna_name, dna_type)
if isinstance(path, tuple) and len(path) > 1 and isinstance(path[-1], int):
# The caller wants to get a single item from an array. The offset we seeked to already
# points to this item. In this case we do not want to look at dna_name.array_size,
# because we want a single item from that array.
return simple_reader(fileobj)
if dna_name.array_size > 1:
return [simple_reader(fileobj) for _ in range(dna_name.array_size)]
return simple_reader(fileobj)

View File

@ -208,7 +208,8 @@ class StructTest(unittest.TestCase):
fileobj = mock.MagicMock(io.BufferedReader)
fileobj.read.return_value = b'\x01\x02\x03\x04\xff\xfe\xfd\xfa\x00dummydata'
val = self.s.field_get(self.FakeHeader(), fileobj, b'path', as_str=False)
val = self.s.field_get(self.FakeHeader(), fileobj, b'path',
as_str=False, null_terminated=True)
self.assertEqual(b'\x01\x02\x03\x04\xff\xfe\xfd\xfa', val)
fileobj.seek.assert_called_with(16, os.SEEK_CUR)
@ -216,8 +217,7 @@ class StructTest(unittest.TestCase):
fileobj = mock.MagicMock(io.BufferedReader)
fileobj.read.return_value = b'\x01\x02\x03\x04\xff\xfe\xfd\xfa\x00dummydata'
val = self.s.field_get(self.FakeHeader(), fileobj, b'path', as_str=False,
nil_terminated=False)
val = self.s.field_get(self.FakeHeader(), fileobj, b'path', as_str=False)
self.assertEqual(b'\x01\x02\x03\x04\xff\xfe\xfd\xfa\x00dummydata', val)
fileobj.seek.assert_called_with(16, os.SEEK_CUR)

View File

@ -28,18 +28,26 @@ class BlendLoadingTest(unittest.TestCase):
# Try low level operation to read a property.
self.bf.fileobj.seek(ob.file_offset, os.SEEK_SET)
loc = ob.dna_type.field_get(self.bf.header, self.bf.fileobj, b'loc')
self.assertEqual(loc, [2.0, 3.0, 5.0])
self.assertEqual([2.0, 3.0, 5.0], loc)
# Try low level operation to read an array element.
self.bf.fileobj.seek(ob.file_offset, os.SEEK_SET)
loc_z = ob.dna_type.field_get(self.bf.header, self.bf.fileobj, (b'loc', 2))
self.assertEqual(5.0, loc_z)
# Try high level operation to read the same property.
loc = ob.get(b'loc')
self.assertEqual(loc, [2.0, 3.0, 5.0])
self.assertEqual([2.0, 3.0, 5.0], loc)
# Try getting a subproperty.
name = ob.get((b'id', b'name'))
self.assertEqual('OBümlaut', name)
loc_z = ob.get((b'loc', 2))
self.assertEqual(5.0, loc_z)
# Try following a pointer.
mesh_ptr = ob.get(b'data')
mesh = self.bf.block_from_offset[mesh_ptr]
mesh = self.bf.block_from_addr[mesh_ptr]
mname = mesh.get((b'id', b'name'))
self.assertEqual('MECube³', mname)