Clarified type refinement, simplified API, custom exception for get_pointer

Type refinement is now only done with BlendFileBlock.refine_type(), and
no longer with sdna_index_refine parameters to various functions. This
simplifies the API at the expense of having to call two simple functions
instead of one more complex one.

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
+    def refine_type_from_index(self, sdna_index: int):
-        assert (type(sdna_index_next) is int)
+        """Change the DNA Struct associated with this block.
-        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(self, dna_type_id):  # TODO(Sybren): port to BAT
+        Use to make a block type more specific, for example when you have a
-        assert (type(dna_type_id) is bytes)
+        modifier but need to access it as SubSurfModifier.
-        self.refine_type_from_index(self.bfile.sdna_index_from_id[dna_type_id])
+
+        :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,
+                   self.get(path_full, default, null_terminated, as_str, base_index))
-                            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,
-    ):
+    ) -> typing.Union[None, 'BlendFileBlock', typing.Any]:
-        if sdna_index_refine is None:
+        """Same as get() but dereferences a pointer.
-            sdna_index_refine = self.sdna_index
-        result = self.get(path, default, sdna_index_refine=sdna_index_refine, base_index=base_index)
 
-        # 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(
+        if result == 0:
-            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:
             return None
 
-    # ----------------------
+        try:
-    # Python convenience API
+            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)
 

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

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)

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'])