diff --git a/blender_asset_tracer/blendfile/__init__.py b/blender_asset_tracer/blendfile/__init__.py index db79676..3c1c354 100644 --- a/blender_asset_tracer/blendfile/__init__.py +++ b/blender_asset_tracer/blendfile/__init__.py @@ -326,41 +326,24 @@ class BlendFileBlock: assert (type(dna_type_id) is bytes) self.refine_type_from_index(self.bfile.sdna_index_from_id[dna_type_id]) - def get_file_offset(self, path, - default=..., - sdna_index_refine=None, - base_index=0, - ): - """ - Return (offset, length) - """ + def get_file_offset(self, path: bytes) -> (int, int): + """Return (offset, length)""" assert isinstance(path, bytes) + # TODO: refactor to just return the length, and check whether this isn't actually + # simply the same as self.size. ofs = self.file_offset - if base_index != 0: - assert (base_index < self.count) - ofs += (self.size // self.count) * base_index - self.bfile.fileobj.seek(ofs, os.SEEK_SET) - - if sdna_index_refine is None: - sdna_index_refine = self.sdna_index - else: - self.bfile.ensure_subtype_smaller(self.sdna_index, sdna_index_refine) - - dna_struct = self.bfile.structs[sdna_index_refine] - field = dna_struct.field_from_path( - self.bfile.header, self.bfile.fileobj, path) - - return self.bfile.fileobj.tell(), field.dna_name.array_size + field, _ = self.dna_type.field_from_path(self.bfile.header.pointer_size, path) + return ofs, field.name.array_size def get(self, path: dna.FieldPath, default=..., sdna_index_refine=None, - null_terminated: typing.Optional[bool]=None, + null_terminated: typing.Optional[bool] = None, as_str=True, base_index=0, - ): + ) -> typing.Any: """Read a property and return the value. :param path: name of the property (like `b'loc'`), tuple of names @@ -395,12 +378,20 @@ class BlendFileBlock: null_terminated=null_terminated, as_str=as_str, ) - def get_recursive_iter(self, path, path_root=b"", + def get_recursive_iter(self, + path: dna.FieldPath, + path_root: dna.FieldPath = b'', default=..., sdna_index_refine=None, - use_nil=True, use_str=True, + null_terminated: typing.Optional[bool] = None, + as_str=True, base_index=0, - ): + ) -> typing.Iterator[typing.Tuple[bytes, typing.Any]]: + """Generator, yields (path, property value) tuples. + + If a property cannot be decoded, a string representing its DNA type + name is used as its value instead, between pointy brackets. + """ if path_root: path_full = ( (path_root if type(path_root) is tuple else (path_root,)) + @@ -409,22 +400,26 @@ class BlendFileBlock: path_full = path try: + # Try accessing as simple property yield (path_full, - self.get(path_full, default, sdna_index_refine, use_nil, use_str, base_index)) - except NotImplementedError as ex: - msg, dna_name, dna_type = ex.args - struct_index = self.bfile.sdna_index_from_id.get(dna_type.dna_type_id, None) + self.get(path_full, default, sdna_index_refine, null_terminated, as_str, + base_index)) + except exceptions.NoReaderImplemented as ex: + # This was not a simple property, so recurse into its DNA Struct. + dna_type = ex.dna_type + struct_index = self.bfile.sdna_index_from_id.get(dna_type.dna_type_id) if struct_index is None: yield (path_full, "<%s>" % dna_type.dna_type_id.decode('ascii')) - else: - struct = self.bfile.structs[struct_index] - for f in struct.fields: - yield from self.get_recursive_iter( - f.dna_name.name_only, path_full, default, None, use_nil, use_str, 0) + return + + # Recurse through the fields. + for f in dna_type.fields: + yield from self.get_recursive_iter(f.name.name_only, path_full, default=default, + null_terminated=null_terminated, as_str=as_str) def items_recursive_iter(self): for k in self.keys(): - yield from self.get_recursive_iter(k, use_str=False) + yield from self.get_recursive_iter(k, as_str=False) def get_data_hash(self): """ diff --git a/blender_asset_tracer/blendfile/dna.py b/blender_asset_tracer/blendfile/dna.py index 2e8bffa..efcd489 100644 --- a/blender_asset_tracer/blendfile/dna.py +++ b/blender_asset_tracer/blendfile/dna.py @@ -2,7 +2,7 @@ import typing import os -from . import dna_io, header +from . import dna_io, header, exceptions # Either a simple path b'propname', or a tuple (b'parentprop', b'actualprop', arrayindex) FieldPath = typing.Union[bytes, typing.Iterable[typing.Union[bytes, int]]] @@ -115,6 +115,14 @@ class Struct: self._fields.append(field) self._fields_by_name[field.name.name_only] = field + @property + def fields(self) -> typing.List[Field]: + """Return the fields of this Struct. + + Do not modify the returned list; use append_field() instead. + """ + return self._fields + def field_from_path(self, pointer_size: int, path: FieldPath) \ @@ -231,8 +239,9 @@ class Struct: 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) + raise exceptions.NoReaderImplemented( + "%r exists but isn't pointer, can't resolve field %r" % (path, dna_name.name_only), + dna_name, dna_type) from None 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 diff --git a/blender_asset_tracer/blendfile/exceptions.py b/blender_asset_tracer/blendfile/exceptions.py index e408b36..8721671 100644 --- a/blender_asset_tracer/blendfile/exceptions.py +++ b/blender_asset_tracer/blendfile/exceptions.py @@ -37,3 +37,18 @@ class BlendFileError(Exception): class NoDNA1Block(BlendFileError): """Raised when the blend file contains no DNA1 block.""" + + +class NoReaderImplemented(NotImplementedError): + """Raised when reading a property of a non-implemented type. + + This indicates that the property should be read using some dna.Struct. + + :type dna_name: blender_asset_tracer.blendfile.dna.Name + :type dna_type: blender_asset_tracer.blendfile.dna.Struct + """ + + def __init__(self, message: str, dna_name, dna_type): + super().__init__(message) + self.dna_name = dna_name + self.dna_type = dna_type diff --git a/tests/test_blendfile_loading.py b/tests/test_blendfile_loading.py index cb24340..e658acc 100644 --- a/tests/test_blendfile_loading.py +++ b/tests/test_blendfile_loading.py @@ -6,7 +6,7 @@ import os from blender_asset_tracer import blendfile -class BlendLoadingTest(unittest.TestCase): +class BlendFileBlockTest(unittest.TestCase): @classmethod def setUpClass(cls): cls.blendfiles = pathlib.Path(__file__).with_name('blendfiles') @@ -18,10 +18,12 @@ class BlendLoadingTest(unittest.TestCase): if self.bf: self.bf.close() - def test_some_properties(self): + def test_loading(self): self.bf = blendfile.BlendFile(self.blendfiles / 'basic_file.blend') self.assertFalse(self.bf.is_compressed) - self.assertEqual(1, len(self.bf.code_index[b'OB'])) + + def test_some_properties(self): + self.bf = blendfile.BlendFile(self.blendfiles / 'basic_file.blend') ob = self.bf.code_index[b'OB'][0] self.assertEqual('Object', ob.dna_type_name) @@ -51,3 +53,29 @@ class BlendLoadingTest(unittest.TestCase): mesh = self.bf.block_from_addr[mesh_ptr] mname = mesh.get((b'id', b'name')) self.assertEqual('MECube³', mname) + + def test_get_recursive_iter(self): + self.bf = blendfile.BlendFile(self.blendfiles / 'basic_file.blend') + ob = self.bf.code_index[b'OB'][0] + assert isinstance(ob, blendfile.BlendFileBlock) + + # No recursing, just an array property. + gen = ob.get_recursive_iter(b'loc') + self.assertEqual([(b'loc', [2.0, 3.0, 5.0])], list(gen)) + + # Recurse into an object + gen = ob.get_recursive_iter(b'id') + self.assertEqual( + [((b'id', b'next'), 0), + ((b'id', b'prev'), 0), + ((b'id', b'newid'), 0), + ((b'id', b'lib'), 0), + ((b'id', b'name'), 'OBümlaut'), + ((b'id', b'flag'), 0), + ((b'id', b'tag'), 1024), + ((b'id', b'us'), 1), + ((b'id', b'icon_id'), 0), + ((b'id', b'recalc'), 0), + ((b'id', b'pad'), 0), + ], + list(gen)[:-2]) # the last 2 properties are pointers and change when saving.