Ported dna.Name, dna.Field, and started with dna.Struct

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)
+                field = dna.Field(dna_type, dna_name, dna_size, dna_offset)
-                dna_struct.fields.append(field)
+                dna_struct.append_field(field)
-                dna_struct.field_from_name[dna_name.name_only] = field
                 dna_offset += dna_size
 
         return structs, sdna_index_from_id

blender_asset_tracer/blendfile/dna.py

@@ -1,20 +1,11 @@
 import os
-import struct
+import typing
 
 
-class DNAName:
+class Name:
-    """
+    """dna.Name is a C-type name stored in the DNA as bytes."""
-    DNAName is a C-type name stored in the DNA
-    """
-    __slots__ = (
-        "name_full",
-        "name_only",
-        "is_pointer",
-        "is_method_pointer",
-        "array_size",
-    )
 
-    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):
+    def as_reference(self, parent) -> bytes:
-        if parent is None:
+        if not parent:
-            result = b''
+            return self.name_only
-        else:
+        return parent + b'.' + self.name_only
-            result = parent + b'.'
 
-        result = result + self.name_only
+    def calc_name_only(self) -> bytes:
-        return result
-
-    def calc_name_only(self):
         result = self.name_full.strip(b'*()')
         index = result.find(b'[')
-        if index != -1:
+        if index == -1:
-            result = result[:index]
             return result
+        return result[:index]
 
-    def calc_is_pointer(self):
+    def calc_is_pointer(self) -> bool:
-        return (b'*' in self.name_full)
+        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
+        partial_name = self.name_full
-        index = temp.find(b'[')
 
-        while index != -1:
+        while True:
-            index_2 = temp.find(b']')
+            idx_start = partial_name.find(b'[')
-            result *= int(temp[index + 1:index_2])
+            if idx_start < 0:
-            temp = temp[index_2 + 1:]
+                break
-            index = temp.find(b'[')
+
+            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:
+class Field:
-    """
+    """dna.Field is a coupled dna.Struct and dna.Name.
-    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",
-    )
 
-    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.name = name
-        self.dna_size = dna_size
+        self.size = size
-        self.dna_offset = dna_offset
+        self.offset = offset
 
 
-class DNAStruct:
+class Struct:
-    """
+    """dna.Struct is a C-type structure stored in the DNA."""
-    DNAStruct is a C-type structure stored in the DNA
-    """
-    __slots__ = (
-        "dna_type_id",
-        "size",
-        "fields",
-        "field_from_name",
-        "user_data",
-    )
 
-    def __init__(self, dna_type_id):
+    def __init__(self, dna_type_id: bytes):
         self.dna_type_id = dna_type_id
-        self.fields = []
+        self._fields = []
-        self.field_from_name = {}
+        self._fields_by_name = {}
-        self.user_data = None
 
     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:
+        if name_tail:
-            handle.seek(field.dna_offset, os.SEEK_CUR)
+            return field.dna_type.field_from_path(pointer_size, name_tail)
-            if index != 0:
+
-                if field.dna_name.is_pointer:
+        offset = field.offset
-                    index_offset = header.pointer_size * index
+        # fileobj.seek(field.offset, os.SEEK_CUR)
+        if index:
+            if field.name.is_pointer:
+                index_offset = pointer_size * index
             else:
                 index_offset = field.dna_type.size * index
-                assert (index_offset < field.dna_size)
+            if index_offset >= field.size:
-                handle.seek(index_offset, os.SEEK_CUR)
+                raise OverflowError('path %r is out of bounds of its DNA type' % path)
-            if not name_tail:  # None or ()
+            # fileobj.seek(index_offset, os.SEEK_CUR)
-                return field
+            offset += index_offset
-            else:
+        return field, offset
-                return field.dna_type.field_from_path(header, handle, name_tail)
 
     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)
-

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')
     USHORT2 = struct.Struct(b'<HH')  # two shorts in a row
@@ -61,6 +61,10 @@ class LittleEndianTypes:
         raise ValueError('unsupported pointer size %d' % fileheader.pointer_size)
 
 
+class LittleEndianTypes(EndianIO):
+    pass
+
+
 class BigEndianTypes(LittleEndianTypes):
     UCHAR = struct.Struct(b'>b')
     USHORT = struct.Struct(b'>H')

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