6bfa4062d7
---
The blendfile module within BAT supports reading data from structs, such
as the Count property on an array modifier. However, blendfile only
supports modifying structs with type "char". This patch adds support for
writing structs of more types in blendfile blocks. Now, writing is
supported for ushort, short, uint, int, float, and ulong types.
The use case that inspired this patch was an instance where a file had
several array modifiers that prevented the file from being opened in
Blender on machines without large amounts of RAM. A solution using the
blendfile module may look like:
```
from blender_asset_tracer import blendfile
from pathlib import Path
b = blendfile.open_cached(Path('flag.blend'), mode='rb+')
for block in b.blocks:
if 'ArrayModifierData' in block.__str__():
try:
print('previous:', block.get(b'count'))
block.set(b'count', 1)
print('current:', block.get(b'count'))
except KeyError:
continue
b.close()
```
This would fail with the exception
`blender_asset_tracer.blendfile.exceptions.NoWriterImplemented: Setting
type Struct(b'int') is not supported for ArrayModifierData.count`. With
this patch, the above code succeeds and the count struct can be set to
a lower number that allows the file to be opened.
This solution implements missing functionality without adding any new
interfaces. A few details are:
* When deciding what type to write to the struct, the value is inferred
from what is given by the caller. If the caller gives a Python int, the
exact type is inferred from the DNA type ID. If they give a float, a
float is written. Otherwise, the existing logic is used to determine
whether to write a string or byte sequence.
* A \_write method was added to dna\_io.py that takes a Python struct
object and a value to write a byte sequence to the file object. This
method is used by public methods appropriately named to indicate what
type they will write.
* The check for whether the caller is trying to write an unsupported
type is left in place, but it has been changed to include types which
are now supported.
* Tests have been added that provide a mock file object, call the new
methods, and confirm that the correct bytes were written.
Reviewed By: sybren
Differential Revision: https://developer.blender.org/D14374
353 lines
12 KiB
Python
353 lines
12 KiB
Python
# ***** BEGIN GPL LICENSE BLOCK *****
|
|
#
|
|
# This program is free software; you can redistribute it and/or
|
|
# modify it under the terms of the GNU General Public License
|
|
# as published by the Free Software Foundation; either version 2
|
|
# of the License, or (at your option) any later version.
|
|
#
|
|
# This program is distributed in the hope that it will be useful,
|
|
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
|
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
|
# GNU General Public License for more details.
|
|
#
|
|
# You should have received a copy of the GNU General Public License
|
|
# along with this program; if not, write to the Free Software Foundation,
|
|
# Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
|
|
#
|
|
# ***** END GPL LICENCE BLOCK *****
|
|
#
|
|
# (c) 2009, At Mind B.V. - Jeroen Bakker
|
|
# (c) 2014, Blender Foundation - Campbell Barton
|
|
# (c) 2018, Blender Foundation - Sybren A. Stüvel
|
|
import logging
|
|
import os
|
|
import typing
|
|
|
|
from . import header, exceptions
|
|
|
|
# Either a simple path b'propname', or a tuple (b'parentprop', b'actualprop', arrayindex)
|
|
FieldPath = typing.Union[bytes, typing.Iterable[typing.Union[bytes, int]]]
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
|
|
class Name:
|
|
"""dna.Name is a C-type name stored in the DNA as bytes."""
|
|
|
|
def __init__(self, name_full: bytes) -> None:
|
|
self.name_full = name_full
|
|
self.name_only = self.calc_name_only()
|
|
self.is_pointer = self.calc_is_pointer()
|
|
self.is_method_pointer = self.calc_is_method_pointer()
|
|
self.array_size = self.calc_array_size()
|
|
|
|
def __repr__(self):
|
|
return "%s(%r)" % (type(self).__qualname__, self.name_full)
|
|
|
|
def as_reference(self, parent) -> bytes:
|
|
if not parent:
|
|
return self.name_only
|
|
return parent + b"." + self.name_only
|
|
|
|
def calc_name_only(self) -> bytes:
|
|
result = self.name_full.strip(b"*()")
|
|
index = result.find(b"[")
|
|
if index == -1:
|
|
return result
|
|
return result[:index]
|
|
|
|
def calc_is_pointer(self) -> bool:
|
|
return b"*" in self.name_full
|
|
|
|
def calc_is_method_pointer(self):
|
|
return b"(*" in self.name_full
|
|
|
|
def calc_array_size(self):
|
|
result = 1
|
|
partial_name = self.name_full
|
|
|
|
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 Field:
|
|
"""dna.Field is a coupled dna.Struct and dna.Name.
|
|
|
|
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) -> None:
|
|
self.dna_type = dna_type
|
|
self.name = name
|
|
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."""
|
|
|
|
log = log.getChild("Struct")
|
|
|
|
def __init__(self, dna_type_id: bytes, size: int = None) -> None:
|
|
"""
|
|
:param dna_type_id: name of the struct in C, like b'AlembicObjectPath'.
|
|
:param size: only for unit tests; typically set after construction by
|
|
BlendFile.decode_structs(). If not set, it is calculated on the fly
|
|
when struct.size is evaluated, based on the available fields.
|
|
"""
|
|
self.dna_type_id = dna_type_id
|
|
self._size = size
|
|
self._fields = [] # type: typing.List[Field]
|
|
self._fields_by_name = {} # type: typing.Dict[bytes, Field]
|
|
|
|
def __repr__(self):
|
|
return "%s(%r)" % (type(self).__qualname__, self.dna_type_id)
|
|
|
|
@property
|
|
def size(self) -> int:
|
|
if self._size is None:
|
|
if not self._fields:
|
|
raise ValueError("Unable to determine size of fieldless %r" % self)
|
|
last_field = max(self._fields, key=lambda f: f.offset)
|
|
self._size = last_field.offset + last_field.size
|
|
return self._size
|
|
|
|
@size.setter
|
|
def size(self, new_size: int):
|
|
self._size = new_size
|
|
|
|
def append_field(self, field: Field):
|
|
self._fields.append(field)
|
|
self._fields_by_name[field.name.name_only] = field
|
|
|
|
@property
|
|
def fields(self) -> typing.List[Field]:
|
|
"""Return the fields of this Struct.
|
|
|
|
Do not modify the returned list; use append_field() instead.
|
|
"""
|
|
return self._fields
|
|
|
|
def has_field(self, field_name: bytes) -> bool:
|
|
return field_name in self._fields_by_name
|
|
|
|
def field_from_path(
|
|
self, pointer_size: int, path: FieldPath
|
|
) -> typing.Tuple[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]' --> (b'array', 4)
|
|
|
|
:returns: the field itself, and its offset taking into account the
|
|
optional index. The offset is relative to the start of the struct,
|
|
i.e. relative to the BlendFileBlock containing the data.
|
|
:raises KeyError: if the field does not exist.
|
|
"""
|
|
if isinstance(path, tuple):
|
|
name = path[0]
|
|
if len(path) >= 2 and not isinstance(path[1], bytes):
|
|
name_tail = path[2:]
|
|
index = path[1]
|
|
assert isinstance(index, int)
|
|
else:
|
|
name_tail = path[1:]
|
|
index = 0
|
|
else:
|
|
name = path
|
|
name_tail = ()
|
|
index = 0
|
|
|
|
if not isinstance(name, bytes):
|
|
raise TypeError("name should be bytes, but is %r" % type(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()))
|
|
)
|
|
|
|
offset = field.offset
|
|
if index:
|
|
if field.name.is_pointer:
|
|
index_offset = pointer_size * index
|
|
else:
|
|
index_offset = field.dna_type.size * index
|
|
if index_offset >= field.size:
|
|
raise OverflowError(
|
|
"path %r is out of bounds of its DNA type %s"
|
|
% (path, field.dna_type)
|
|
)
|
|
offset += index_offset
|
|
|
|
if name_tail:
|
|
subval, suboff = field.dna_type.field_from_path(pointer_size, name_tail)
|
|
return subval, suboff + offset
|
|
|
|
return field, offset
|
|
|
|
def field_get(
|
|
self,
|
|
file_header: header.BlendFileHeader,
|
|
fileobj: typing.IO[bytes],
|
|
path: FieldPath,
|
|
default=...,
|
|
null_terminated=True,
|
|
as_str=True,
|
|
) -> typing.Tuple[typing.Optional[Field], typing.Any]:
|
|
"""Read the value of the field from the blend file.
|
|
|
|
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. Be careful with this
|
|
default when reading binary data.
|
|
:param as_str: When True, automatically decode bytes to string
|
|
(assumes UTF-8 encoding).
|
|
:returns: The field instance and the value. If a default value was passed
|
|
and the field was not found, (None, default) is returned.
|
|
"""
|
|
try:
|
|
field, offset = self.field_from_path(file_header.pointer_size, path)
|
|
except KeyError:
|
|
if default is ...:
|
|
raise
|
|
return None, default
|
|
|
|
fileobj.seek(offset, os.SEEK_CUR)
|
|
|
|
dna_type = field.dna_type
|
|
dna_name = field.name
|
|
endian = file_header.endian
|
|
|
|
# Some special cases (pointers, strings/bytes)
|
|
if dna_name.is_pointer:
|
|
return field, endian.read_pointer(fileobj, file_header.pointer_size)
|
|
if dna_type.dna_type_id == b"char":
|
|
return field, self._field_get_char(
|
|
file_header, fileobj, field, null_terminated, as_str
|
|
)
|
|
|
|
simple_readers = {
|
|
b"int": endian.read_int,
|
|
b"short": endian.read_short,
|
|
b"uint64_t": endian.read_ulong,
|
|
b"float": endian.read_float,
|
|
}
|
|
try:
|
|
simple_reader = simple_readers[dna_type.dna_type_id]
|
|
except KeyError:
|
|
raise exceptions.NoReaderImplemented(
|
|
"%r exists but not simple type (%r), can't resolve field %r"
|
|
% (path, dna_type.dna_type_id.decode(), dna_name.name_only),
|
|
dna_name,
|
|
dna_type,
|
|
) from None
|
|
|
|
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 field, simple_reader(fileobj)
|
|
|
|
if dna_name.array_size > 1:
|
|
return field, [simple_reader(fileobj) for _ in range(dna_name.array_size)]
|
|
return field, simple_reader(fileobj)
|
|
|
|
def _field_get_char(
|
|
self,
|
|
file_header: header.BlendFileHeader,
|
|
fileobj: typing.IO[bytes],
|
|
field: "Field",
|
|
null_terminated: typing.Optional[bool],
|
|
as_str: bool,
|
|
) -> typing.Any:
|
|
dna_name = field.name
|
|
endian = file_header.endian
|
|
|
|
if field.size == 1:
|
|
# Single char, assume it's bitflag or int value, and not a string/bytes data...
|
|
return endian.read_char(fileobj)
|
|
|
|
if null_terminated or (null_terminated is None and as_str):
|
|
data = endian.read_bytes0(fileobj, dna_name.array_size)
|
|
else:
|
|
data = fileobj.read(dna_name.array_size)
|
|
|
|
if as_str:
|
|
return data.decode("utf8")
|
|
return data
|
|
|
|
def field_set(
|
|
self,
|
|
file_header: header.BlendFileHeader,
|
|
fileobj: typing.IO[bytes],
|
|
path: bytes,
|
|
value: typing.Any,
|
|
):
|
|
"""Write a value to the blend file.
|
|
|
|
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).
|
|
"""
|
|
assert isinstance(path, bytes), "path should be bytes, but is %s" % type(path)
|
|
|
|
field, offset = self.field_from_path(file_header.pointer_size, path)
|
|
|
|
dna_type = field.dna_type
|
|
dna_name = field.name
|
|
endian = file_header.endian
|
|
|
|
if dna_type.dna_type_id not in endian.accepted_types():
|
|
msg = "Setting type %r is not supported for %s.%s" % (
|
|
dna_type,
|
|
self.dna_type_id.decode(),
|
|
dna_name.name_full.decode(),
|
|
)
|
|
raise exceptions.NoWriterImplemented(msg, dna_name, dna_type)
|
|
|
|
fileobj.seek(offset, os.SEEK_CUR)
|
|
|
|
if self.log.isEnabledFor(logging.DEBUG):
|
|
filepos = fileobj.tell()
|
|
if isinstance(value, (int, float)):
|
|
thing = dna_type.dna_type_id.decode()
|
|
else:
|
|
thing = "string" if isinstance(value, str) else "bytes"
|
|
self.log.debug(
|
|
"writing %s %r at file offset %d / %x", thing, value, filepos, filepos
|
|
)
|
|
|
|
if isinstance(value, (int, float)):
|
|
return endian.accepted_types()[dna_type.dna_type_id](fileobj, value)
|
|
elif isinstance(value, str):
|
|
return endian.write_string(fileobj, value, dna_name.array_size)
|
|
else:
|
|
return endian.write_bytes(fileobj, value, dna_name.array_size)
|