Fixed issue reading array items and renamed some stuff

blender_asset_tracer/blendfile/__init__.py

@@ -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,8 +123,7 @@ 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_addr = {block.addr_old: block for block in self.blocks
-        self.block_from_offset = {block.addr_old: block for block in self.blocks
                                 if block.code != b'ENDB'}
 
     def __repr__(self):
@@ -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)

blender_asset_tracer/blendfile/dna.py

@@ -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)

tests/test_blendfile_dna.py

@@ -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,
+        val = self.s.field_get(self.FakeHeader(), fileobj, b'path', as_str=False)
-                               nil_terminated=False)
         self.assertEqual(b'\x01\x02\x03\x04\xff\xfe\xfd\xfa\x00dummydata', val)
         fileobj.seek.assert_called_with(16, os.SEEK_CUR)
 

tests/test_blendfile_loading.py

@@ -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)