diff --git a/blender_asset_tracer/blendfile/__init__.py b/blender_asset_tracer/blendfile/__init__.py index 91c041b..db27905 100644 --- a/blender_asset_tracer/blendfile/__init__.py +++ b/blender_asset_tracer/blendfile/__init__.py @@ -315,15 +315,30 @@ class BlendFileBlock: def dna_type_name(self) -> str: return self.dna_type.dna_type_id.decode('ascii') - def refine_type_from_index(self, sdna_index_next): # TODO(Sybren): port to BAT - assert (type(sdna_index_next) is int) - sdna_index_curr = self.sdna_index - self.bfile.ensure_subtype_smaller(sdna_index_curr, sdna_index_next) - self.sdna_index = sdna_index_next + def refine_type_from_index(self, sdna_index: int): + """Change the DNA Struct associated with this block. - def refine_type(self, dna_type_id): # TODO(Sybren): port to BAT - assert (type(dna_type_id) is bytes) - self.refine_type_from_index(self.bfile.sdna_index_from_id[dna_type_id]) + Use to make a block type more specific, for example when you have a + modifier but need to access it as SubSurfModifier. + + :param sdna_index: the SDNA index of the DNA type. + """ + assert type(sdna_index) is int + sdna_index_curr = self.sdna_index + self.bfile.ensure_subtype_smaller(sdna_index_curr, sdna_index) + self.sdna_index = sdna_index + + def refine_type(self, dna_type_id: bytes): + """Change the DNA Struct associated with this block. + + Use to make a block type more specific, for example when you have a + modifier but need to access it as SubSurfModifier. + + :param dna_type_id: the name of the DNA type. + """ + assert isinstance(dna_type_id, bytes) + sdna_index = self.bfile.sdna_index_from_id[dna_type_id] + self.refine_type_from_index(sdna_index) def get_file_offset(self, path: bytes) -> (int, int): # TODO(Sybren): port to BAT """Return (offset, length)""" @@ -338,7 +353,6 @@ class BlendFileBlock: def get(self, path: dna.FieldPath, default=..., - sdna_index_refine=None, null_terminated=True, as_str=True, base_index=0, @@ -365,30 +379,17 @@ class BlendFileBlock: ofs += (self.size // self.count) * base_index self.bfile.fileobj.seek(ofs, os.SEEK_SET) - dna_struct = self._get_struct(sdna_index_refine) + dna_struct = self.bfile.structs[self.sdna_index] return dna_struct.field_get( self.bfile.header, self.bfile.fileobj, path, default=default, null_terminated=null_terminated, as_str=as_str, ) - def _get_struct(self, sdna_index_refine) -> dna.Struct: - """Gets the (possibly refined) struct for this block.""" - - if sdna_index_refine is None: - index = self.sdna_index - else: - self.bfile.ensure_subtype_smaller(self.sdna_index, sdna_index_refine) - index = sdna_index_refine - - dna_struct = self.bfile.structs[index] - return dna_struct - def get_recursive_iter(self, path: dna.FieldPath, path_root: dna.FieldPath = b'', default=..., - sdna_index_refine=None, null_terminated=True, as_str=True, base_index=0, @@ -408,8 +409,7 @@ class BlendFileBlock: try: # Try accessing as simple property yield (path_full, - self.get(path_full, default, sdna_index_refine, null_terminated, as_str, - base_index)) + self.get(path_full, default, 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 @@ -451,38 +451,31 @@ class BlendFileBlock: return dna_struct.field_set( self.bfile.header, self.bfile.handle, path, value) - # --------------- - # Utility get/set - # - # avoid inline pointer casting def get_pointer( - self, path, + self, path: dna.FieldPath, default=..., - sdna_index_refine=None, base_index=0, - ): - if sdna_index_refine is None: - sdna_index_refine = self.sdna_index - result = self.get(path, default, sdna_index_refine=sdna_index_refine, base_index=base_index) + ) -> typing.Union[None, 'BlendFileBlock', typing.Any]: + """Same as get() but dereferences a pointer. - # default + :raises exceptions.SegmentationFault: when there is no datablock with + the pointed-to address. + """ + result = self.get(path, default=default, base_index=base_index) + + # If it's not an integer, we have no pointer to follow and this may + # actually be a non-pointer property. if type(result) is not int: return result - assert (self.bfile.structs[sdna_index_refine].field_from_path( - self.bfile.header, self.bfile.handle, path).dna_name.is_pointer) - if result != 0: - # possible (but unlikely) - # that this fails and returns None - # maybe we want to raise some exception in this case - return self.bfile.find_block_from_offset(result) - else: + if result == 0: return None - # ---------------------- - # Python convenience API + try: + return self.bfile.block_from_addr[result] + except KeyError: + raise exceptions.SegmentationFault('address does not exist', path, result) - # dict like access def __getitem__(self, path: dna.FieldPath): return self.get(path, as_str=False) diff --git a/blender_asset_tracer/blendfile/exceptions.py b/blender_asset_tracer/blendfile/exceptions.py index 8721671..e23e566 100644 --- a/blender_asset_tracer/blendfile/exceptions.py +++ b/blender_asset_tracer/blendfile/exceptions.py @@ -52,3 +52,12 @@ class NoReaderImplemented(NotImplementedError): super().__init__(message) self.dna_name = dna_name self.dna_type = dna_type + + +class SegmentationFault(Exception): + """Raised when a pointer to a non-existant datablock was dereferenced.""" + + def __init__(self, message: str, field_path, address: int): + super().__init__(message) + self.field_path = field_path + self.address = address diff --git a/blender_asset_tracer/blendfile/iterators.py b/blender_asset_tracer/blendfile/iterators.py new file mode 100644 index 0000000..3e0c775 --- /dev/null +++ b/blender_asset_tracer/blendfile/iterators.py @@ -0,0 +1,9 @@ +from . import BlendFileBlock + + +def listbase(block: BlendFileBlock) -> BlendFileBlock: + """Generator, yields all blocks in the ListBase linked list.""" + while block: + yield block + next_ptr = block[b'next'] + block = block.bfile.find_block_from_address(next_ptr) diff --git a/tests/blendfiles/with_sequencer.blend b/tests/blendfiles/with_sequencer.blend new file mode 100644 index 0000000..31701d7 Binary files /dev/null and b/tests/blendfiles/with_sequencer.blend differ diff --git a/tests/test_blendfile_loading.py b/tests/test_blendfile_loading.py index 9415cbe..039eb33 100644 --- a/tests/test_blendfile_loading.py +++ b/tests/test_blendfile_loading.py @@ -4,20 +4,26 @@ import unittest import os from blender_asset_tracer import blendfile +from blender_asset_tracer.blendfile import iterators -class BlendFileBlockTest(unittest.TestCase): +class AbstractBlendFileTest(unittest.TestCase): @classmethod def setUpClass(cls): cls.blendfiles = pathlib.Path(__file__).with_name('blendfiles') def setUp(self): - self.bf = blendfile.BlendFile(self.blendfiles / 'basic_file.blend') + self.bf = None def tearDown(self): if self.bf: self.bf.close() + +class BlendFileBlockTest(AbstractBlendFileTest): + def setUp(self): + self.bf = blendfile.BlendFile(self.blendfiles / 'basic_file.blend') + def test_loading(self): self.assertFalse(self.bf.is_compressed) @@ -148,3 +154,60 @@ class BlendFileBlockTest(unittest.TestCase): ob = self.bf.code_index[b'OB'][0] assert isinstance(ob, blendfile.BlendFileBlock) self.assertEqual('OBümlaut', ob[b'id', b'name'].decode()) + + +class PointerTest(AbstractBlendFileTest): + def setUp(self): + self.bf = blendfile.BlendFile(self.blendfiles / 'with_sequencer.blend') + + def test_get_pointer_and_listbase(self): + scenes = self.bf.code_index[b'SC'] + self.assertEqual(1, len(scenes), 'expecting 1 scene') + scene = scenes[0] + self.assertEqual(b'SCScene', scene[b'id', b'name']) + + ed_ptr = scene[b'ed'] + self.assertEqual(140051431100936, ed_ptr) + + ed = scene.get_pointer(b'ed') + self.assertEqual(140051431100936, ed.addr_old) + + seqbase = ed.get_pointer((b'seqbase', b'first')) + self.assertIsNotNone(seqbase) + + types = { + b'SQBlack': 28, + b'SQCross': 8, + b'SQPink': 28, + } + seq = None + for seq in iterators.listbase(seqbase): + seq.refine_type(b'Sequence') + name = seq[b'name'] + expected_type = types[name] + self.assertEqual(expected_type, seq[b'type']) + + # The last 'seq' from the loop should be the last in the list. + seq_next = seq.get_pointer(b'next') + self.assertIsNone(seq_next) + + def test_refine_sdna_by_name(self): + scene = self.bf.code_index[b'SC'][0] + ed = scene.get_pointer(b'ed') + + seq = ed.get_pointer((b'seqbase', b'first')) + + # This is very clear to me: + seq.refine_type(b'Sequence') + self.assertEqual(b'SQBlack', seq[b'name']) + self.assertEqual(28, seq[b'type']) + + def test_refine_sdna_by_idx(self): + scene = self.bf.code_index[b'SC'][0] + ed = scene.get_pointer(b'ed') + seq = ed.get_pointer((b'seqbase', b'first')) + + sdna_idx_sequence = self.bf.sdna_index_from_id[b'Sequence'] + seq.refine_type_from_index(sdna_idx_sequence) + self.assertEqual(b'SQBlack', seq[b'name']) + self.assertEqual(28, seq[b'type'])