diff --git a/blender_asset_tracer/blendfile/__init__.py b/blender_asset_tracer/blendfile/__init__.py index 2482e04..8697721 100644 --- a/blender_asset_tracer/blendfile/__init__.py +++ b/blender_asset_tracer/blendfile/__init__.py @@ -102,18 +102,19 @@ class BlendFile: self.sdna_index_from_id = {} self.block_from_offset = {} + self.load_dna1_block() + def load_dna1_block(self): """Read the blend file to load its DNA structure to memory.""" - fileobj = self.fileobj while True: - block = BlendFileBlock(fileobj, self) + block = BlendFileBlock(self) if block.code == b'ENDB': break if block.code == b'DNA1': self.structs, self.sdna_index_from_id = self.decode_structs(block) else: - fileobj.seek(block.size, os.SEEK_CUR) + self.fileobj.seek(block.size, os.SEEK_CUR) self.blocks.append(block) self.code_index[block.code].append(block) @@ -255,8 +256,8 @@ class BlendFileBlock: old_structure = struct.Struct(b'4sI') """old blend files ENDB block structure""" - def __init__(self, fileobj: typing.BinaryIO, bfile: BlendFile): - self.file = bfile + def __init__(self, bfile: BlendFile): + self.bfile = bfile # Defaults; actual values are set by interpreting the block header. self.code = b'' @@ -267,11 +268,12 @@ class BlendFileBlock: self.file_offset = 0 """Offset in bytes from start of file to beginning of the data block. - Points to the data after the block header. - """ + Points to the data after the block header. + """ + self.endian = bfile.header.endian header_struct = bfile.block_header_struct - data = fileobj.read(header_struct.size) + data = bfile.fileobj.read(header_struct.size) if len(data) != header_struct.size: self.log.warning("Blend file %s seems to be truncated, " "expected %d bytes but could read only %d", @@ -286,17 +288,17 @@ class BlendFileBlock: if len(data) <= 15: self.log.debug('interpreting block as old-style ENB block') blockheader = self.old_structure.unpack(data) - self.code = dna_io.read_data0(blockheader[0]) + self.code = self.endian.read_data0(blockheader[0]) return blockheader = header_struct.unpack(data) - self.code = dna_io.read_data0(blockheader[0]) + self.code = self.endian.read_data0(blockheader[0]) if self.code != b'ENDB': self.size = blockheader[1] self.addr_old = blockheader[2] self.sdna_index = blockheader[3] self.count = blockheader[4] - self.file_offset = fileobj.tell() + self.file_offset = bfile.fileobj.tell() def __str__(self): return "<%s.%s (%s), size=%d at %s>" % ( @@ -308,22 +310,22 @@ class BlendFileBlock: ) @property - def dna_type(self): - return self.file.structs[self.sdna_index] + def dna_type(self) -> dna.Struct: + return self.bfile.structs[self.sdna_index] @property - def dna_type_name(self): + def dna_type_name(self) -> str: return self.dna_type.dna_type_id.decode('ascii') def refine_type_from_index(self, sdna_index_next): assert (type(sdna_index_next) is int) sdna_index_curr = self.sdna_index - self.file.ensure_subtype_smaller(sdna_index_curr, sdna_index_next) + self.bfile.ensure_subtype_smaller(sdna_index_curr, sdna_index_next) self.sdna_index = sdna_index_next def refine_type(self, dna_type_id): assert (type(dna_type_id) is bytes) - self.refine_type_from_index(self.file.sdna_index_from_id[dna_type_id]) + self.refine_type_from_index(self.bfile.sdna_index_from_id[dna_type_id]) def get_file_offset(self, path, default=..., @@ -333,24 +335,24 @@ class BlendFileBlock: """ Return (offset, length) """ - assert (type(path) is bytes) + assert isinstance(path, bytes) ofs = self.file_offset if base_index != 0: assert (base_index < self.count) ofs += (self.size // self.count) * base_index - self.file.handle.seek(ofs, os.SEEK_SET) + self.bfile.fileobj.seek(ofs, os.SEEK_SET) if sdna_index_refine is None: sdna_index_refine = self.sdna_index else: - self.file.ensure_subtype_smaller(self.sdna_index, sdna_index_refine) + self.bfile.ensure_subtype_smaller(self.sdna_index, sdna_index_refine) - dna_struct = self.file.structs[sdna_index_refine] + dna_struct = self.bfile.structs[sdna_index_refine] field = dna_struct.field_from_path( - self.file.header, self.file.handle, path) + self.bfile.header, self.bfile.fileobj, path) - return (self.file.handle.tell(), field.dna_name.array_size) + return self.bfile.fileobj.tell(), field.dna_name.array_size def get(self, path, default=..., @@ -363,16 +365,16 @@ class BlendFileBlock: if base_index != 0: assert (base_index < self.count) ofs += (self.size // self.count) * base_index - self.file.handle.seek(ofs, os.SEEK_SET) + self.bfile.fileobj.seek(ofs, os.SEEK_SET) if sdna_index_refine is None: sdna_index_refine = self.sdna_index else: - self.file.ensure_subtype_smaller(self.sdna_index, sdna_index_refine) + self.bfile.ensure_subtype_smaller(self.sdna_index, sdna_index_refine) - dna_struct = self.file.structs[sdna_index_refine] + dna_struct = self.bfile.structs[sdna_index_refine] return dna_struct.field_get( - self.file.header, self.file.handle, path, + self.bfile.header, self.bfile.fileobj, path, default=default, nil_terminated=use_nil, as_str=use_str, ) @@ -395,11 +397,11 @@ class BlendFileBlock: 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.file.sdna_index_from_id.get(dna_type.dna_type_id, None) + struct_index = self.bfile.sdna_index_from_id.get(dna_type.dna_type_id, None) if struct_index is None: yield (path_full, "<%s>" % dna_type.dna_type_id.decode('ascii')) else: - struct = self.file.structs[struct_index] + 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) @@ -433,13 +435,13 @@ class BlendFileBlock: if sdna_index_refine is None: sdna_index_refine = self.sdna_index else: - self.file.ensure_subtype_smaller(self.sdna_index, sdna_index_refine) + self.bfile.ensure_subtype_smaller(self.sdna_index, sdna_index_refine) - dna_struct = self.file.structs[sdna_index_refine] - self.file.handle.seek(self.file_offset, os.SEEK_SET) - self.file.is_modified = True + dna_struct = self.bfile.structs[sdna_index_refine] + self.bfile.handle.seek(self.file_offset, os.SEEK_SET) + self.bfile.is_modified = True return dna_struct.field_set( - self.file.header, self.file.handle, path, value) + self.bfile.header, self.bfile.handle, path, value) # --------------- # Utility get/set @@ -459,13 +461,13 @@ class BlendFileBlock: if type(result) is not int: return result - assert (self.file.structs[sdna_index_refine].field_from_path( - self.file.header, self.file.handle, path).dna_name.is_pointer) + 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.file.find_block_from_offset(result) + return self.bfile.find_block_from_offset(result) else: return None diff --git a/blender_asset_tracer/blendfile/dna.py b/blender_asset_tracer/blendfile/dna.py index b1dfefd..96285b6 100644 --- a/blender_asset_tracer/blendfile/dna.py +++ b/blender_asset_tracer/blendfile/dna.py @@ -76,6 +76,9 @@ class Field: self.size = size self.offset = offset + def __repr__(self): + return '<%r %r (%s)>' % (type(self).__qualname__, self.name, self.dna_type) + class Struct: """dna.Struct is a C-type structure stored in the DNA.""" diff --git a/blender_asset_tracer/blendfile/dna_io.py b/blender_asset_tracer/blendfile/dna_io.py index 78cde8b..5b5c56c 100644 --- a/blender_asset_tracer/blendfile/dna_io.py +++ b/blender_asset_tracer/blendfile/dna_io.py @@ -17,7 +17,10 @@ class EndianIO: @classmethod def _read(cls, fileobj: typing.BinaryIO, typestruct: struct.Struct): data = fileobj.read(typestruct.size) - return typestruct.unpack(data)[0] + try: + return typestruct.unpack(data)[0] + except struct.error as ex: + raise struct.error('%s (read %d bytes)' % (ex, len(data))) from None @classmethod def read_char(cls, fileobj: typing.BinaryIO): diff --git a/blender_asset_tracer/blendfile/header.py b/blender_asset_tracer/blendfile/header.py index b9a67b2..b195f51 100644 --- a/blender_asset_tracer/blendfile/header.py +++ b/blender_asset_tracer/blendfile/header.py @@ -1,4 +1,5 @@ import logging +import os import pathlib import struct import typing @@ -12,13 +13,14 @@ class BlendFileHeader: """ BlendFileHeader represents the first 12 bytes of a blend file. - it contains information about the hardware architecture, which is relevant + It contains information about the hardware architecture, which is relevant to the structure of the rest of the file. """ structure = struct.Struct(b'7s1s1s3s') def __init__(self, fileobj: typing.BinaryIO, path: pathlib.Path): log.debug("reading blend-file-header %s", path) + fileobj.seek(0, os.SEEK_SET) header = fileobj.read(self.structure.size) values = self.structure.unpack(header) diff --git a/tests/blendfiles/basic_file.blend b/tests/blendfiles/basic_file.blend new file mode 100644 index 0000000..b573b7b Binary files /dev/null and b/tests/blendfiles/basic_file.blend differ diff --git a/tests/test_blendfile_loading.py b/tests/test_blendfile_loading.py new file mode 100644 index 0000000..2020c75 --- /dev/null +++ b/tests/test_blendfile_loading.py @@ -0,0 +1,45 @@ +import pathlib +import unittest + +import os + +from blender_asset_tracer import blendfile + + +class BlendLoadingTest(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.blendfiles = pathlib.Path(__file__).with_name('blendfiles') + + def setUp(self): + self.bf = None + + def tearDown(self): + if self.bf: + self.bf.close() + + def test_some_properties(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'])) + ob = self.bf.code_index[b'OB'][0] + self.assertEqual('Object', ob.dna_type_name) + + # 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]) + + # Try high level operation to read the same property. + loc = ob.get(b'loc') + self.assertEqual(loc, [2.0, 3.0, 5.0]) + + # Try getting a subproperty. + name = ob.get((b'id', b'name')) + self.assertEqual('OBümlaut', name) + + # Try following a pointer. + mesh_ptr = ob.get(b'data') + mesh = self.bf.block_from_offset[mesh_ptr] + mname = mesh.get((b'id', b'name')) + self.assertEqual('MECube³', mname)