diff --git a/blender_asset_tracer/blendfile/__init__.py b/blender_asset_tracer/blendfile/__init__.py index f0461e4..4f09351 100644 --- a/blender_asset_tracer/blendfile/__init__.py +++ b/blender_asset_tracer/blendfile/__init__.py @@ -29,7 +29,7 @@ import pathlib import tempfile import typing -from . import exceptions, dna_io +from . import exceptions, dna_io, dna log = logging.getLogger(__name__) @@ -64,7 +64,7 @@ class BlendFile: """ self.filepath = path - fileobj = path.open(mode) + fileobj = path.open(mode, buffering=FILE_BUFFER_SIZE) magic = fileobj.read(len(BLENDFILE_MAGIC)) if magic == BLENDFILE_MAGIC: @@ -190,7 +190,7 @@ class BlendFile: for _ in range(names_len): typename = dna_io.read_data0_offset(data, offset) offset = offset + len(typename) + 1 - typenames.append(DNAName(typename)) + typenames.append(dna.Name(typename)) offset = pad_up_4(offset) offset += 4 @@ -199,7 +199,7 @@ class BlendFile: self.log.debug("building #%d types" % types_len) for _ in range(types_len): dna_type_id = dna_io.read_data0_offset(data, offset) - types.append(DNAStruct(dna_type_id)) + types.append(dna.Struct(dna_type_id)) offset += len(dna_type_id) + 1 offset = pad_up_4(offset) @@ -239,9 +239,8 @@ class BlendFile: else: dna_size = dna_type.size * dna_name.array_size - field = DNAField(dna_type, dna_name, dna_size, dna_offset) - dna_struct.fields.append(field) - dna_struct.field_from_name[dna_name.name_only] = field + field = dna.Field(dna_type, dna_name, dna_size, dna_offset) + dna_struct.append_field(field) dna_offset += dna_size return structs, sdna_index_from_id diff --git a/blender_asset_tracer/blendfile/dna.py b/blender_asset_tracer/blendfile/dna.py index a438a28..e0c4796 100644 --- a/blender_asset_tracer/blendfile/dna.py +++ b/blender_asset_tracer/blendfile/dna.py @@ -1,20 +1,11 @@ import os -import struct +import typing -class DNAName: - """ - DNAName is a C-type name stored in the DNA - """ - __slots__ = ( - "name_full", - "name_only", - "is_pointer", - "is_method_pointer", - "array_size", - ) +class Name: + """dna.Name is a C-type name stored in the DNA as bytes.""" - def __init__(self, name_full): + def __init__(self, name_full: bytes): self.name_full = name_full self.name_only = self.calc_name_only() self.is_pointer = self.calc_is_pointer() @@ -24,100 +15,95 @@ class DNAName: def __repr__(self): return '%s(%r)' % (type(self).__qualname__, self.name_full) - def as_reference(self, parent): - if parent is None: - result = b'' - else: - result = parent + b'.' + def as_reference(self, parent) -> bytes: + if not parent: + return self.name_only + return parent + b'.' + self.name_only - result = result + self.name_only - return result - - def calc_name_only(self): + def calc_name_only(self) -> bytes: result = self.name_full.strip(b'*()') index = result.find(b'[') - if index != -1: - result = result[:index] - return result + if index == -1: + return result + return result[:index] - def calc_is_pointer(self): - return (b'*' in self.name_full) + def calc_is_pointer(self) -> bool: + return b'*' in self.name_full def calc_is_method_pointer(self): - return (b'(*' in self.name_full) + return b'(*' in self.name_full def calc_array_size(self): result = 1 - temp = self.name_full - index = temp.find(b'[') + partial_name = self.name_full - while index != -1: - index_2 = temp.find(b']') - result *= int(temp[index + 1:index_2]) - temp = temp[index_2 + 1:] - index = temp.find(b'[') + while True: + idx_start = partial_name.find(b'[') + if idx_start < 0: + break + + idx_stop = partial_name.find(b']') + result *= int(partial_name[idx_start + 1:idx_stop]) + partial_name = partial_name[idx_stop + 1:] return result -class DNAField: - """ - DNAField is a coupled DNAStruct and DNAName - and cache offset for reuse - """ - __slots__ = ( - # DNAName - "dna_name", - # tuple of 3 items - # [bytes (struct name), int (struct size), DNAStruct] - "dna_type", - # size on-disk - "dna_size", - # cached info (avoid looping over fields each time) - "dna_offset", - ) +class Field: + """dna.Field is a coupled dna.Struct and dna.Name. - def __init__(self, dna_type, dna_name, dna_size, dna_offset): + It also contains the file offset in bytes. + + :ivar name: the name of the field. + :ivar dna_type: the type of the field. + :ivar size: size of the field on disk, in bytes. + :ivar offset: cached offset of the field, in bytes. + """ + + def __init__(self, + dna_type: 'Struct', + name: Name, + size: int, + offset: int): self.dna_type = dna_type - self.dna_name = dna_name - self.dna_size = dna_size - self.dna_offset = dna_offset + self.name = name + self.size = size + self.offset = offset -class DNAStruct: - """ - DNAStruct is a C-type structure stored in the DNA - """ - __slots__ = ( - "dna_type_id", - "size", - "fields", - "field_from_name", - "user_data", - ) +class Struct: + """dna.Struct is a C-type structure stored in the DNA.""" - def __init__(self, dna_type_id): + def __init__(self, dna_type_id: bytes): self.dna_type_id = dna_type_id - self.fields = [] - self.field_from_name = {} - self.user_data = None + self._fields = [] + self._fields_by_name = {} def __repr__(self): return '%s(%r)' % (type(self).__qualname__, self.dna_type_id) - def field_from_path(self, header, handle, path): + def append_field(self, field: Field): + self._fields.append(field) + self._fields_by_name[field.name.name_only] = field + + def field_from_path(self, + pointer_size: int, + path: typing.Union[bytes, typing.Iterable[typing.Union[bytes, int]]]) \ + -> typing.Tuple[typing.Optional[Field], int]: """ Support lookups as bytes or a tuple of bytes and optional index. C style 'id.name' --> (b'id', b'name') - C style 'array[4]' --> ('array', 4) + C style 'array[4]' --> (b'array', 4) + + :returns: the field itself, and its offset taking into account the optional index. """ - if type(path) is tuple: + if isinstance(path, (tuple, list)): name = path[0] - if len(path) >= 2 and type(path[1]) is not bytes: + if len(path) >= 2 and not isinstance(path[1], bytes): name_tail = path[2:] index = path[1] - assert (type(index) is int) + assert isinstance(index, int) else: name_tail = path[1:] index = 0 @@ -126,23 +112,29 @@ class DNAStruct: name_tail = None index = 0 - assert (type(name) is bytes) + if not isinstance(name, bytes): + raise TypeError('name should be bytes, but is %r' % type(name)) - field = self.field_from_name.get(name) + field = self._fields_by_name.get(name) + if not field: + raise KeyError('%r has no field %r, only %r' % + (self, name, sorted(self._fields_by_name.keys()))) - if field is not None: - handle.seek(field.dna_offset, os.SEEK_CUR) - if index != 0: - if field.dna_name.is_pointer: - index_offset = header.pointer_size * index - else: - index_offset = field.dna_type.size * index - assert (index_offset < field.dna_size) - handle.seek(index_offset, os.SEEK_CUR) - if not name_tail: # None or () - return field + if name_tail: + return field.dna_type.field_from_path(pointer_size, name_tail) + + offset = field.offset + # fileobj.seek(field.offset, os.SEEK_CUR) + if index: + if field.name.is_pointer: + index_offset = pointer_size * index else: - return field.dna_type.field_from_path(header, handle, name_tail) + index_offset = field.dna_type.size * index + if index_offset >= field.size: + raise OverflowError('path %r is out of bounds of its DNA type' % path) + # fileobj.seek(index_offset, os.SEEK_CUR) + offset += index_offset + return field, offset def field_get(self, header, handle, path, default=..., @@ -155,7 +147,7 @@ class DNAStruct: else: raise KeyError("%r not found in %r (%r)" % ( - path, [f.dna_name.name_only for f in self.fields], + path, [f.dna_name.name_only for f in self._fields], self.dna_type_id)) dna_type = field.dna_type @@ -217,4 +209,3 @@ class DNAStruct: else: raise NotImplementedError("Setting %r is not yet supported for %r" % (dna_type, dna_name), dna_name, dna_type) - diff --git a/blender_asset_tracer/blendfile/dna_io.py b/blender_asset_tracer/blendfile/dna_io.py index f76054f..36d7e77 100644 --- a/blender_asset_tracer/blendfile/dna_io.py +++ b/blender_asset_tracer/blendfile/dna_io.py @@ -4,7 +4,7 @@ import struct import typing -class LittleEndianTypes: +class EndianIO: UCHAR = struct.Struct(b'b') USHORT = struct.Struct(b'>H') diff --git a/tests/test_blendfile_dna.py b/tests/test_blendfile_dna.py new file mode 100644 index 0000000..f1bc874 --- /dev/null +++ b/tests/test_blendfile_dna.py @@ -0,0 +1,85 @@ +import unittest + +from blender_asset_tracer.blendfile import dna + + +class NameTest(unittest.TestCase): + def test_simple_name(self): + n = dna.Name(b'Suzanne') + self.assertEqual(n.name_full, b'Suzanne') + self.assertEqual(n.name_only, b'Suzanne') + self.assertFalse(n.is_pointer) + self.assertFalse(n.is_method_pointer) + self.assertEqual(n.array_size, 1) + + def test_pointer(self): + n = dna.Name(b'*marker') + self.assertEqual(n.name_full, b'*marker') + self.assertEqual(n.name_only, b'marker') + self.assertTrue(n.is_pointer) + self.assertFalse(n.is_method_pointer) + self.assertEqual(n.array_size, 1) + + def test_method_pointer(self): + n = dna.Name(b'(*delta_cache)()') + self.assertEqual(n.name_full, b'(*delta_cache)()') + self.assertEqual(n.name_only, b'delta_cache') + self.assertTrue(n.is_pointer) + self.assertTrue(n.is_method_pointer) + self.assertEqual(n.array_size, 1) + + def test_simple_array(self): + n = dna.Name(b'flame_smoke_color[3]') + self.assertEqual(n.name_full, b'flame_smoke_color[3]') + self.assertEqual(n.name_only, b'flame_smoke_color') + self.assertFalse(n.is_pointer) + self.assertFalse(n.is_method_pointer) + self.assertEqual(n.array_size, 3) + + def test_nested_array(self): + n = dna.Name(b'pattern_corners[4][2]') + self.assertEqual(n.name_full, b'pattern_corners[4][2]') + self.assertEqual(n.name_only, b'pattern_corners') + self.assertFalse(n.is_pointer) + self.assertFalse(n.is_method_pointer) + self.assertEqual(n.array_size, 8) + + def test_pointer_array(self): + n = dna.Name(b'*mtex[18]') + self.assertEqual(n.name_full, b'*mtex[18]') + self.assertEqual(n.name_only, b'mtex') + self.assertTrue(n.is_pointer) + self.assertFalse(n.is_method_pointer) + self.assertEqual(n.array_size, 18) + + def test_repr(self): + self.assertEqual(repr(dna.Name(b'Suzanne')), "Name(b'Suzanne')") + self.assertEqual(repr(dna.Name(b'*marker')), "Name(b'*marker')") + self.assertEqual(repr(dna.Name(b'(*delta_cache)()')), "Name(b'(*delta_cache)()')") + self.assertEqual(repr(dna.Name(b'flame_smoke_color[3]')), "Name(b'flame_smoke_color[3]')") + self.assertEqual(repr(dna.Name(b'pattern_corners[4][2]')), "Name(b'pattern_corners[4][2]')") + self.assertEqual(repr(dna.Name(b'*mtex[18]')), "Name(b'*mtex[18]')") + + def test_as_reference(self): + n = dna.Name(b'(*delta_cache)()') + self.assertEqual(n.as_reference(None), b'delta_cache') + self.assertEqual(n.as_reference(b''), b'delta_cache') + self.assertEqual(n.as_reference(b'parent'), b'parent.delta_cache') + + +class StructTest(unittest.TestCase): + def test_field_from_path(self): + s = dna.Struct(b'AlembicObjectPath') + f_next = dna.Field(s, dna.Name(b'*next'), 8, 0) + f_prev = dna.Field(s, dna.Name(b'*prev'), 8, 8) + f_path = dna.Field(dna.Struct(b'char'), dna.Name(b'path[4096]'), 4096, 16) + f_pointer = dna.Field(dna.Struct(b'char'), dna.Name(b'*ptr'), 3 * 8, 16 + 4096) + s.append_field(f_next) + s.append_field(f_prev) + s.append_field(f_path) + s.append_field(f_pointer) + + psize = 8 + self.assertEqual(s.field_from_path(psize, b'path'), (f_path, 16)) + self.assertEqual(s.field_from_path(psize, (b'prev', b'path')), (f_path, 16)) + self.assertEqual(s.field_from_path(psize, (b'ptr', 2)), (f_pointer, 16 + 4096 + 2 * psize))