Allow block.get() to return the dna.Field

This is needed by the upcoming dependency tracer.
This commit is contained in:
Sybren A. Stüvel 2018-02-26 18:15:14 +01:00
parent a56e985cdc
commit 59c0b6df4c
4 changed files with 29 additions and 21 deletions

View File

@ -386,6 +386,7 @@ class BlendFileBlock:
null_terminated=True, null_terminated=True,
as_str=False, as_str=False,
base_index=0, base_index=0,
return_field=False
) -> typing.Any: ) -> typing.Any:
"""Read a property and return the value. """Read a property and return the value.
@ -400,6 +401,8 @@ class BlendFileBlock:
when reading binary data. when reading binary data.
:param as_str: When True, automatically decode bytes to string :param as_str: When True, automatically decode bytes to string
(assumes UTF-8 encoding). (assumes UTF-8 encoding).
:param return_field: When True, returns tuple (dna.Field, value).
Otherwise just returns the value.
""" """
ofs = self.file_offset ofs = self.file_offset
if base_index != 0: if base_index != 0:
@ -410,11 +413,14 @@ class BlendFileBlock:
self.bfile.fileobj.seek(ofs, os.SEEK_SET) self.bfile.fileobj.seek(ofs, os.SEEK_SET)
dna_struct = self.bfile.structs[self.sdna_index] dna_struct = self.bfile.structs[self.sdna_index]
return dna_struct.field_get( field, value = dna_struct.field_get(
self.bfile.header, self.bfile.fileobj, path, self.bfile.header, self.bfile.fileobj, path,
default=default, default=default,
null_terminated=null_terminated, as_str=as_str, null_terminated=null_terminated, as_str=as_str,
) )
if return_field:
return value, field
return value
def get_recursive_iter(self, def get_recursive_iter(self,
path: dna.FieldPath, path: dna.FieldPath,

View File

@ -184,7 +184,7 @@ class Struct:
default=..., default=...,
null_terminated=True, null_terminated=True,
as_str=True, as_str=True,
): ) -> typing.Tuple[typing.Optional[Field], typing.Any]:
"""Read the value of the field from the blend file. """Read the value of the field from the blend file.
Assumes the file pointer of `fileobj` is seek()ed to the start of the Assumes the file pointer of `fileobj` is seek()ed to the start of the
@ -201,13 +201,15 @@ class Struct:
default when reading binary data. default when reading binary data.
:param as_str: When True, automatically decode bytes to string :param as_str: When True, automatically decode bytes to string
(assumes UTF-8 encoding). (assumes UTF-8 encoding).
:returns: The field instance and the value. If a default value was passed
and the field was not found, (None, default) is returned.
""" """
try: try:
field, offset = self.field_from_path(file_header.pointer_size, path) field, offset = self.field_from_path(file_header.pointer_size, path)
except KeyError: except KeyError:
if default is ...: if default is ...:
raise raise
return default return None, default
fileobj.seek(offset, os.SEEK_CUR) fileobj.seek(offset, os.SEEK_CUR)
@ -217,19 +219,19 @@ class Struct:
# Some special cases (pointers, strings/bytes) # Some special cases (pointers, strings/bytes)
if dna_name.is_pointer: if dna_name.is_pointer:
return endian.read_pointer(fileobj, file_header.pointer_size) return field, endian.read_pointer(fileobj, file_header.pointer_size)
if dna_type.dna_type_id == b'char': if dna_type.dna_type_id == b'char':
if field.size == 1: if field.size == 1:
# Single char, assume it's bitflag or int value, and not a string/bytes data... # Single char, assume it's bitflag or int value, and not a string/bytes data...
return endian.read_char(fileobj) return field, endian.read_char(fileobj)
if null_terminated or (null_terminated is None and as_str): if null_terminated or (null_terminated is None and as_str):
data = endian.read_bytes0(fileobj, dna_name.array_size) data = endian.read_bytes0(fileobj, dna_name.array_size)
else: else:
data = fileobj.read(dna_name.array_size) data = fileobj.read(dna_name.array_size)
if as_str: if as_str:
return data.decode('utf8') return field, data.decode('utf8')
return data return field, data
simple_readers = { simple_readers = {
b'int': endian.read_int, b'int': endian.read_int,
@ -249,11 +251,11 @@ class Struct:
# The caller wants to get a single item from an array. The offset we seeked to already # 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, # 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. # because we want a single item from that array.
return simple_reader(fileobj) return field, simple_reader(fileobj)
if dna_name.array_size > 1: if dna_name.array_size > 1:
return [simple_reader(fileobj) for _ in range(dna_name.array_size)] return field, [simple_reader(fileobj) for _ in range(dna_name.array_size)]
return simple_reader(fileobj) return field, simple_reader(fileobj)
def field_set(self, def field_set(self,
file_header: header.BlendFileHeader, file_header: header.BlendFileHeader,

View File

@ -144,7 +144,7 @@ class StructTest(unittest.TestCase):
def test_simple_field_get(self): def test_simple_field_get(self):
fileobj = mock.MagicMock(io.BufferedReader) fileobj = mock.MagicMock(io.BufferedReader)
fileobj.read.return_value = b'\x01\x02\x03\x04\xff\xfe\xfd\xfa' fileobj.read.return_value = b'\x01\x02\x03\x04\xff\xfe\xfd\xfa'
val = self.s.field_get(self.FakeHeader(), fileobj, b'numbah') _, val = self.s.field_get(self.FakeHeader(), fileobj, b'numbah')
self.assertEqual(val, 0x1020304fffefdfa) self.assertEqual(val, 0x1020304fffefdfa)
fileobj.seek.assert_called_with(4136, os.SEEK_CUR) fileobj.seek.assert_called_with(4136, os.SEEK_CUR)
@ -152,7 +152,7 @@ class StructTest(unittest.TestCase):
def test_field_get_default(self): def test_field_get_default(self):
fileobj = mock.MagicMock(io.BufferedReader) fileobj = mock.MagicMock(io.BufferedReader)
fileobj.read.side_effect = RuntimeError fileobj.read.side_effect = RuntimeError
val = self.s.field_get(self.FakeHeader(), fileobj, b'nonexistant', default=519871531) _, val = self.s.field_get(self.FakeHeader(), fileobj, b'nonexistant', default=519871531)
self.assertEqual(val, 519871531) self.assertEqual(val, 519871531)
fileobj.seek.assert_not_called() fileobj.seek.assert_not_called()
@ -176,7 +176,7 @@ class StructTest(unittest.TestCase):
def test_pointer_field_get(self): def test_pointer_field_get(self):
fileobj = mock.MagicMock(io.BufferedReader) fileobj = mock.MagicMock(io.BufferedReader)
fileobj.read.return_value = b'\xf0\x9f\xa6\x87\x00dum' fileobj.read.return_value = b'\xf0\x9f\xa6\x87\x00dum'
val = self.s.field_get(self.FakeHeader(), fileobj, b'ptr') _, val = self.s.field_get(self.FakeHeader(), fileobj, b'ptr')
self.assertEqual(0xf09fa6870064756d, val) self.assertEqual(0xf09fa6870064756d, val)
fileobj.seek.assert_called_with(4112, os.SEEK_CUR) fileobj.seek.assert_called_with(4112, os.SEEK_CUR)
@ -184,7 +184,7 @@ class StructTest(unittest.TestCase):
def test_string_field_get(self): def test_string_field_get(self):
fileobj = mock.MagicMock(io.BufferedReader) fileobj = mock.MagicMock(io.BufferedReader)
fileobj.read.return_value = b'\xf0\x9f\xa6\x87\x00dummydata' fileobj.read.return_value = b'\xf0\x9f\xa6\x87\x00dummydata'
val = self.s.field_get(self.FakeHeader(), fileobj, b'path', as_str=True) _, val = self.s.field_get(self.FakeHeader(), fileobj, b'path', as_str=True)
self.assertEqual('🦇', val) self.assertEqual('🦇', val)
fileobj.seek.assert_called_with(16, os.SEEK_CUR) fileobj.seek.assert_called_with(16, os.SEEK_CUR)
@ -192,7 +192,7 @@ class StructTest(unittest.TestCase):
def test_string_field_get_single_char(self): def test_string_field_get_single_char(self):
fileobj = mock.MagicMock(io.BufferedReader) fileobj = mock.MagicMock(io.BufferedReader)
fileobj.read.return_value = b'\xf0' fileobj.read.return_value = b'\xf0'
val = self.s.field_get(self.FakeHeader(), fileobj, b'bitflag') _, val = self.s.field_get(self.FakeHeader(), fileobj, b'bitflag')
self.assertEqual(0xf0, val) self.assertEqual(0xf0, val)
fileobj.seek.assert_called_with(4152, os.SEEK_CUR) fileobj.seek.assert_called_with(4152, os.SEEK_CUR)
@ -208,7 +208,7 @@ class StructTest(unittest.TestCase):
fileobj = mock.MagicMock(io.BufferedReader) fileobj = mock.MagicMock(io.BufferedReader)
fileobj.read.return_value = b'\x01\x02\x03\x04\xff\xfe\xfd\xfa\x00dummydata' 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)
self.assertEqual(b'\x01\x02\x03\x04\xff\xfe\xfd\xfa', val) self.assertEqual(b'\x01\x02\x03\x04\xff\xfe\xfd\xfa', val)
fileobj.seek.assert_called_with(16, os.SEEK_CUR) fileobj.seek.assert_called_with(16, os.SEEK_CUR)
@ -216,8 +216,8 @@ class StructTest(unittest.TestCase):
fileobj = mock.MagicMock(io.BufferedReader) fileobj = mock.MagicMock(io.BufferedReader)
fileobj.read.return_value = b'\x01\x02\x03\x04\xff\xfe\xfd\xfa\x00dummydata' fileobj.read.return_value = b'\x01\x02\x03\x04\xff\xfe\xfd\xfa\x00dummydata'
val = self.s.field_get(self.FakeHeader(), fileobj, b'path', _, val = self.s.field_get(self.FakeHeader(), fileobj, b'path',
as_str=False, null_terminated=False) as_str=False, null_terminated=False)
self.assertEqual(b'\x01\x02\x03\x04\xff\xfe\xfd\xfa\x00dummydata', val) self.assertEqual(b'\x01\x02\x03\x04\xff\xfe\xfd\xfa\x00dummydata', val)
fileobj.seek.assert_called_with(16, os.SEEK_CUR) fileobj.seek.assert_called_with(16, os.SEEK_CUR)
@ -225,7 +225,7 @@ class StructTest(unittest.TestCase):
fileobj = mock.MagicMock(io.BufferedReader) fileobj = mock.MagicMock(io.BufferedReader)
fileobj.read.side_effect = (b'@333', b'@2\x8f\\') fileobj.read.side_effect = (b'@333', b'@2\x8f\\')
val = self.s.field_get(self.FakeHeader(), fileobj, b'floaty') _, val = self.s.field_get(self.FakeHeader(), fileobj, b'floaty')
self.assertAlmostEqual(2.8, val[0]) self.assertAlmostEqual(2.8, val[0])
self.assertAlmostEqual(2.79, val[1]) self.assertAlmostEqual(2.79, val[1])
fileobj.seek.assert_called_with(4144, os.SEEK_CUR) fileobj.seek.assert_called_with(4144, os.SEEK_CUR)

View File

@ -19,12 +19,12 @@ class BlendFileBlockTest(AbstractBlendFileTest):
# Try low level operation to read a property. # Try low level operation to read a property.
self.bf.fileobj.seek(ob.file_offset, os.SEEK_SET) self.bf.fileobj.seek(ob.file_offset, os.SEEK_SET)
loc = ob.dna_type.field_get(self.bf.header, self.bf.fileobj, b'loc') _, loc = ob.dna_type.field_get(self.bf.header, self.bf.fileobj, b'loc')
self.assertEqual([2.0, 3.0, 5.0], loc) self.assertEqual([2.0, 3.0, 5.0], loc)
# Try low level operation to read an array element. # Try low level operation to read an array element.
self.bf.fileobj.seek(ob.file_offset, os.SEEK_SET) 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)) _, loc_z = ob.dna_type.field_get(self.bf.header, self.bf.fileobj, (b'loc', 2))
self.assertEqual(5.0, loc_z) self.assertEqual(5.0, loc_z)
# Try high level operation to read the same property. # Try high level operation to read the same property.