Clarified type refinement, simplified API, custom exception for get_pointer
Type refinement is now only done with BlendFileBlock.refine_type(), and no longer with sdna_index_refine parameters to various functions. This simplifies the API at the expense of having to call two simple functions instead of one more complex one.
This commit is contained in:
Sybren A. Stüvel
parent
9c8766a28c
commit
934a8e210e
@ -315,15 +315,30 @@ class BlendFileBlock:
|
||||
def dna_type_name(self) -> str:
|
||||
return self.dna_type.dna_type_id.decode('ascii')
|
||||
|
||||
def refine_type_from_index(self, sdna_index_next): # TODO(Sybren): port to BAT
|
||||
assert (type(sdna_index_next) is int)
|
||||
sdna_index_curr = self.sdna_index
|
||||
self.bfile.ensure_subtype_smaller(sdna_index_curr, sdna_index_next)
|
||||
self.sdna_index = sdna_index_next
|
||||
def refine_type_from_index(self, sdna_index: int):
|
||||
"""Change the DNA Struct associated with this block.
|
||||
|
||||
def refine_type(self, dna_type_id): # TODO(Sybren): port to BAT
|
||||
assert (type(dna_type_id) is bytes)
|
||||
self.refine_type_from_index(self.bfile.sdna_index_from_id[dna_type_id])
|
||||
Use to make a block type more specific, for example when you have a
|
||||
modifier but need to access it as SubSurfModifier.
|
||||
|
||||
:param sdna_index: the SDNA index of the DNA type.
|
||||
"""
|
||||
assert type(sdna_index) is int
|
||||
sdna_index_curr = self.sdna_index
|
||||
self.bfile.ensure_subtype_smaller(sdna_index_curr, sdna_index)
|
||||
self.sdna_index = sdna_index
|
||||
|
||||
def refine_type(self, dna_type_id: bytes):
|
||||
"""Change the DNA Struct associated with this block.
|
||||
|
||||
Use to make a block type more specific, for example when you have a
|
||||
modifier but need to access it as SubSurfModifier.
|
||||
|
||||
:param dna_type_id: the name of the DNA type.
|
||||
"""
|
||||
assert isinstance(dna_type_id, bytes)
|
||||
sdna_index = self.bfile.sdna_index_from_id[dna_type_id]
|
||||
self.refine_type_from_index(sdna_index)
|
||||
|
||||
def get_file_offset(self, path: bytes) -> (int, int): # TODO(Sybren): port to BAT
|
||||
"""Return (offset, length)"""
|
||||
@ -338,7 +353,6 @@ class BlendFileBlock:
|
||||
def get(self,
|
||||
path: dna.FieldPath,
|
||||
default=...,
|
||||
sdna_index_refine=None,
|
||||
null_terminated=True,
|
||||
as_str=True,
|
||||
base_index=0,
|
||||
@ -365,30 +379,17 @@ class BlendFileBlock:
|
||||
ofs += (self.size // self.count) * base_index
|
||||
self.bfile.fileobj.seek(ofs, os.SEEK_SET)
|
||||
|
||||
dna_struct = self._get_struct(sdna_index_refine)
|
||||
dna_struct = self.bfile.structs[self.sdna_index]
|
||||
return dna_struct.field_get(
|
||||
self.bfile.header, self.bfile.fileobj, path,
|
||||
default=default,
|
||||
null_terminated=null_terminated, as_str=as_str,
|
||||
)
|
||||
|
||||
def _get_struct(self, sdna_index_refine) -> dna.Struct:
|
||||
"""Gets the (possibly refined) struct for this block."""
|
||||
|
||||
if sdna_index_refine is None:
|
||||
index = self.sdna_index
|
||||
else:
|
||||
self.bfile.ensure_subtype_smaller(self.sdna_index, sdna_index_refine)
|
||||
index = sdna_index_refine
|
||||
|
||||
dna_struct = self.bfile.structs[index]
|
||||
return dna_struct
|
||||
|
||||
def get_recursive_iter(self,
|
||||
path: dna.FieldPath,
|
||||
path_root: dna.FieldPath = b'',
|
||||
default=...,
|
||||
sdna_index_refine=None,
|
||||
null_terminated=True,
|
||||
as_str=True,
|
||||
base_index=0,
|
||||
@ -408,8 +409,7 @@ class BlendFileBlock:
|
||||
try:
|
||||
# Try accessing as simple property
|
||||
yield (path_full,
|
||||
self.get(path_full, default, sdna_index_refine, null_terminated, as_str,
|
||||
base_index))
|
||||
self.get(path_full, default, null_terminated, as_str, base_index))
|
||||
except exceptions.NoReaderImplemented as ex:
|
||||
# This was not a simple property, so recurse into its DNA Struct.
|
||||
dna_type = ex.dna_type
|
||||
@ -451,38 +451,31 @@ class BlendFileBlock:
|
||||
return dna_struct.field_set(
|
||||
self.bfile.header, self.bfile.handle, path, value)
|
||||
|
||||
# ---------------
|
||||
# Utility get/set
|
||||
#
|
||||
# avoid inline pointer casting
|
||||
def get_pointer(
|
||||
self, path,
|
||||
self, path: dna.FieldPath,
|
||||
default=...,
|
||||
sdna_index_refine=None,
|
||||
base_index=0,
|
||||
):
|
||||
if sdna_index_refine is None:
|
||||
sdna_index_refine = self.sdna_index
|
||||
result = self.get(path, default, sdna_index_refine=sdna_index_refine, base_index=base_index)
|
||||
) -> typing.Union[None, 'BlendFileBlock', typing.Any]:
|
||||
"""Same as get() but dereferences a pointer.
|
||||
|
||||
# default
|
||||
:raises exceptions.SegmentationFault: when there is no datablock with
|
||||
the pointed-to address.
|
||||
"""
|
||||
result = self.get(path, default=default, base_index=base_index)
|
||||
|
||||
# If it's not an integer, we have no pointer to follow and this may
|
||||
# actually be a non-pointer property.
|
||||
if type(result) is not int:
|
||||
return result
|
||||
|
||||
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.bfile.find_block_from_offset(result)
|
||||
else:
|
||||
if result == 0:
|
||||
return None
|
||||
|
||||
# ----------------------
|
||||
# Python convenience API
|
||||
try:
|
||||
return self.bfile.block_from_addr[result]
|
||||
except KeyError:
|
||||
raise exceptions.SegmentationFault('address does not exist', path, result)
|
||||
|
||||
# dict like access
|
||||
def __getitem__(self, path: dna.FieldPath):
|
||||
return self.get(path, as_str=False)
|
||||
|
||||
|
||||
@ -52,3 +52,12 @@ class NoReaderImplemented(NotImplementedError):
|
||||
super().__init__(message)
|
||||
self.dna_name = dna_name
|
||||
self.dna_type = dna_type
|
||||
|
||||
|
||||
class SegmentationFault(Exception):
|
||||
"""Raised when a pointer to a non-existant datablock was dereferenced."""
|
||||
|
||||
def __init__(self, message: str, field_path, address: int):
|
||||
super().__init__(message)
|
||||
self.field_path = field_path
|
||||
self.address = address
|
||||
|
||||
9
blender_asset_tracer/blendfile/iterators.py
Normal file
9
blender_asset_tracer/blendfile/iterators.py
Normal file
@ -0,0 +1,9 @@
|
||||
from . import BlendFileBlock
|
||||
|
||||
|
||||
def listbase(block: BlendFileBlock) -> BlendFileBlock:
|
||||
"""Generator, yields all blocks in the ListBase linked list."""
|
||||
while block:
|
||||
yield block
|
||||
next_ptr = block[b'next']
|
||||
block = block.bfile.find_block_from_address(next_ptr)
|
||||
BIN
tests/blendfiles/with_sequencer.blend
Normal file
BIN
tests/blendfiles/with_sequencer.blend
Normal file
Binary file not shown.
@ -4,20 +4,26 @@ import unittest
|
||||
import os
|
||||
|
||||
from blender_asset_tracer import blendfile
|
||||
from blender_asset_tracer.blendfile import iterators
|
||||
|
||||
|
||||
class BlendFileBlockTest(unittest.TestCase):
|
||||
class AbstractBlendFileTest(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.blendfiles = pathlib.Path(__file__).with_name('blendfiles')
|
||||
|
||||
def setUp(self):
|
||||
self.bf = blendfile.BlendFile(self.blendfiles / 'basic_file.blend')
|
||||
self.bf = None
|
||||
|
||||
def tearDown(self):
|
||||
if self.bf:
|
||||
self.bf.close()
|
||||
|
||||
|
||||
class BlendFileBlockTest(AbstractBlendFileTest):
|
||||
def setUp(self):
|
||||
self.bf = blendfile.BlendFile(self.blendfiles / 'basic_file.blend')
|
||||
|
||||
def test_loading(self):
|
||||
self.assertFalse(self.bf.is_compressed)
|
||||
|
||||
@ -148,3 +154,60 @@ class BlendFileBlockTest(unittest.TestCase):
|
||||
ob = self.bf.code_index[b'OB'][0]
|
||||
assert isinstance(ob, blendfile.BlendFileBlock)
|
||||
self.assertEqual('OBümlaut', ob[b'id', b'name'].decode())
|
||||
|
||||
|
||||
class PointerTest(AbstractBlendFileTest):
|
||||
def setUp(self):
|
||||
self.bf = blendfile.BlendFile(self.blendfiles / 'with_sequencer.blend')
|
||||
|
||||
def test_get_pointer_and_listbase(self):
|
||||
scenes = self.bf.code_index[b'SC']
|
||||
self.assertEqual(1, len(scenes), 'expecting 1 scene')
|
||||
scene = scenes[0]
|
||||
self.assertEqual(b'SCScene', scene[b'id', b'name'])
|
||||
|
||||
ed_ptr = scene[b'ed']
|
||||
self.assertEqual(140051431100936, ed_ptr)
|
||||
|
||||
ed = scene.get_pointer(b'ed')
|
||||
self.assertEqual(140051431100936, ed.addr_old)
|
||||
|
||||
seqbase = ed.get_pointer((b'seqbase', b'first'))
|
||||
self.assertIsNotNone(seqbase)
|
||||
|
||||
types = {
|
||||
b'SQBlack': 28,
|
||||
b'SQCross': 8,
|
||||
b'SQPink': 28,
|
||||
}
|
||||
seq = None
|
||||
for seq in iterators.listbase(seqbase):
|
||||
seq.refine_type(b'Sequence')
|
||||
name = seq[b'name']
|
||||
expected_type = types[name]
|
||||
self.assertEqual(expected_type, seq[b'type'])
|
||||
|
||||
# The last 'seq' from the loop should be the last in the list.
|
||||
seq_next = seq.get_pointer(b'next')
|
||||
self.assertIsNone(seq_next)
|
||||
|
||||
def test_refine_sdna_by_name(self):
|
||||
scene = self.bf.code_index[b'SC'][0]
|
||||
ed = scene.get_pointer(b'ed')
|
||||
|
||||
seq = ed.get_pointer((b'seqbase', b'first'))
|
||||
|
||||
# This is very clear to me:
|
||||
seq.refine_type(b'Sequence')
|
||||
self.assertEqual(b'SQBlack', seq[b'name'])
|
||||
self.assertEqual(28, seq[b'type'])
|
||||
|
||||
def test_refine_sdna_by_idx(self):
|
||||
scene = self.bf.code_index[b'SC'][0]
|
||||
ed = scene.get_pointer(b'ed')
|
||||
seq = ed.get_pointer((b'seqbase', b'first'))
|
||||
|
||||
sdna_idx_sequence = self.bf.sdna_index_from_id[b'Sequence']
|
||||
seq.refine_type_from_index(sdna_idx_sequence)
|
||||
self.assertEqual(b'SQBlack', seq[b'name'])
|
||||
self.assertEqual(28, seq[b'type'])
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user