diff --git a/blender_asset_tracer/blendfile/__init__.py b/blender_asset_tracer/blendfile/__init__.py index 8697721..db79676 100644 --- a/blender_asset_tracer/blendfile/__init__.py +++ b/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,9 +123,8 @@ 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_offset = {block.addr_old: block for block in self.blocks - if block.code != b'ENDB'} + self.block_from_addr = {block.addr_old: block for block in self.blocks + if block.code != b'ENDB'} def __repr__(self): clsname = self.__class__.__qualname__ @@ -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) diff --git a/blender_asset_tracer/blendfile/dna.py b/blender_asset_tracer/blendfile/dna.py index 96285b6..2e8bffa 100644 --- a/blender_asset_tracer/blendfile/dna.py +++ b/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) diff --git a/tests/test_blendfile_dna.py b/tests/test_blendfile_dna.py index f848fbc..126bff4 100644 --- a/tests/test_blendfile_dna.py +++ b/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, - nil_terminated=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\x00dummydata', val) fileobj.seek.assert_called_with(16, os.SEEK_CUR) diff --git a/tests/test_blendfile_loading.py b/tests/test_blendfile_loading.py index 2020c75..cb24340 100644 --- a/tests/test_blendfile_loading.py +++ b/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)