From 6bfa4062d74908e196dfa613c99810d9161e2234 Mon Sep 17 00:00:00 2001 From: William Harrell Date: Fri, 25 Mar 2022 11:58:43 +0100 Subject: [PATCH] Support for int, float types in BlendFileBlock.set --- 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 --- CHANGELOG.md | 1 + blender_asset_tracer/blendfile/dna.py | 11 ++- blender_asset_tracer/blendfile/dna_io.py | 58 +++++++++++++ tests/test_blendfile_dna.py | 102 +++++++++++++++++++++++ 4 files changed, 169 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 04fe48e..5004574 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,7 @@ changed functionality, fixed bugs). # Version 1.12 (2022-03-25) - Removed "late imports", to help isolate Blender add-ons bundling BAT from each other. +- Support writing `int` and `float` types. - Decided to *not* support the Shaman API of Flamenco 3.x in BAT. The support for that protocol will be implemented in the Flamenco 3.x add-on for Blender, and not in BAT itself. A new future version of BAT will remove support for the Shaman API altogether. # Version 1.11 (2022-02-18) diff --git a/blender_asset_tracer/blendfile/dna.py b/blender_asset_tracer/blendfile/dna.py index 116476e..061c894 100644 --- a/blender_asset_tracer/blendfile/dna.py +++ b/blender_asset_tracer/blendfile/dna.py @@ -324,7 +324,7 @@ class Struct: dna_name = field.name endian = file_header.endian - if dna_type.dna_type_id != b"char": + 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(), @@ -336,12 +336,17 @@ class Struct: if self.log.isEnabledFor(logging.DEBUG): filepos = fileobj.tell() - thing = "string" if isinstance(value, str) else "bytes" + 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, str): + 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) diff --git a/blender_asset_tracer/blendfile/dna_io.py b/blender_asset_tracer/blendfile/dna_io.py index d7c3e70..5596449 100644 --- a/blender_asset_tracer/blendfile/dna_io.py +++ b/blender_asset_tracer/blendfile/dna_io.py @@ -44,34 +44,70 @@ class EndianIO: except struct.error as ex: raise struct.error("%s (read %d bytes)" % (ex, len(data))) from None + @classmethod + def _write(cls, fileobj: typing.IO[bytes], typestruct: struct.Struct, value: typing.Any): + try: + data = typestruct.pack(value) + except struct.error as ex: + raise struct.error(f"{ex} (write '{value}')") + return fileobj.write(data) + @classmethod def read_char(cls, fileobj: typing.IO[bytes]): return cls._read(fileobj, cls.UCHAR) + @classmethod + def write_char(cls, fileobj: typing.IO[bytes], value: int): + return cls._write(fileobj, cls.UCHAR, value) + @classmethod def read_ushort(cls, fileobj: typing.IO[bytes]): return cls._read(fileobj, cls.USHORT) + @classmethod + def write_ushort(cls, fileobj: typing.IO[bytes], value: int): + return cls._write(fileobj, cls.USHORT, value) + @classmethod def read_short(cls, fileobj: typing.IO[bytes]): return cls._read(fileobj, cls.SSHORT) + @classmethod + def write_short(cls, fileobj: typing.IO[bytes], value: int): + return cls._write(fileobj, cls.SSHORT, value) + @classmethod def read_uint(cls, fileobj: typing.IO[bytes]): return cls._read(fileobj, cls.UINT) + @classmethod + def write_uint(cls, fileobj: typing.IO[bytes], value: int): + return cls._write(fileobj, cls.UINT, value) + @classmethod def read_int(cls, fileobj: typing.IO[bytes]): return cls._read(fileobj, cls.SINT) + @classmethod + def write_int(cls, fileobj: typing.IO[bytes], value: int): + return cls._write(fileobj, cls.SINT, value) + @classmethod def read_float(cls, fileobj: typing.IO[bytes]): return cls._read(fileobj, cls.FLOAT) + @classmethod + def write_float(cls, fileobj: typing.IO[bytes], value: float): + return cls._write(fileobj, cls.FLOAT, value) + @classmethod def read_ulong(cls, fileobj: typing.IO[bytes]): return cls._read(fileobj, cls.ULONG) + @classmethod + def write_ulong(cls, fileobj: typing.IO[bytes], value: int): + return cls._write(fileobj, cls.ULONG, value) + @classmethod def read_pointer(cls, fileobj: typing.IO[bytes], pointer_size: int): """Read a pointer from a file.""" @@ -82,6 +118,16 @@ class EndianIO: return cls.read_ulong(fileobj) raise ValueError("unsupported pointer size %d" % pointer_size) + @classmethod + def write_pointer(cls, fileobj: typing.IO[bytes], pointer_size: int, value: int): + """Write a pointer to a file.""" + + if pointer_size == 4: + return cls.write_uint(fileobj, value) + if pointer_size == 8: + return cls.write_ulong(fileobj, value) + raise ValueError("unsupported pointer size %d" % pointer_size) + @classmethod def write_string( cls, fileobj: typing.IO[bytes], astring: str, fieldlen: int @@ -149,6 +195,18 @@ class EndianIO: return data return data[:add] + @classmethod + def accepted_types(cls): + return { + b"char": cls.write_char, + b"ushort": cls.write_ushort, + b"short": cls.write_short, + b"uint": cls.write_uint, + b"int": cls.write_int, + b"ulong": cls.write_ulong, + b"float": cls.write_float + } + class LittleEndianTypes(EndianIO): pass diff --git a/tests/test_blendfile_dna.py b/tests/test_blendfile_dna.py index 5de5dbd..31a3719 100644 --- a/tests/test_blendfile_dna.py +++ b/tests/test_blendfile_dna.py @@ -1,5 +1,6 @@ import io import os +import struct import unittest from unittest import mock @@ -84,7 +85,12 @@ class StructTest(unittest.TestCase): def setUp(self): self.s = dna.Struct(b"AlembicObjectPath") self.s_char = dna.Struct(b"char", 1) + self.s_ushort = dna.Struct(b"ushort", 2) + self.s_short = dna.Struct(b"short", 2) + self.s_uint = dna.Struct(b"uint", 4) + self.s_int = dna.Struct(b"int", 4) self.s_float = dna.Struct(b"float", 4) + self.s_ulong = dna.Struct(b"ulong", 8) self.s_uint64 = dna.Struct(b"uint64_t", 8) self.s_uint128 = dna.Struct(b"uint128_t", 16) # non-supported type @@ -96,6 +102,13 @@ class StructTest(unittest.TestCase): self.f_floaty = dna.Field(self.s_float, dna.Name(b"floaty[2]"), 2 * 4, 4144) self.f_flag = dna.Field(self.s_char, dna.Name(b"bitflag"), 1, 4152) self.f_bignum = dna.Field(self.s_uint128, dna.Name(b"bignum"), 16, 4153) + self.f_testchar = dna.Field(self.s_char, dna.Name(b"testchar"), 1, 4169) + self.f_testushort = dna.Field(self.s_ushort, dna.Name(b"testushort"), 2, 4170) + self.f_testshort = dna.Field(self.s_short, dna.Name(b"testshort"), 2, 4172) + self.f_testuint = dna.Field(self.s_uint, dna.Name(b"testuint"), 4, 4174) + self.f_testint = dna.Field(self.s_int, dna.Name(b"testint"), 4, 4178) + self.f_testfloat = dna.Field(self.s_float, dna.Name(b"testfloat"), 4, 4182) + self.f_testulong = dna.Field(self.s_ulong, dna.Name(b"testulong"), 8, 4186) self.s.append_field(self.f_next) self.s.append_field(self.f_prev) @@ -105,6 +118,13 @@ class StructTest(unittest.TestCase): self.s.append_field(self.f_floaty) self.s.append_field(self.f_flag) self.s.append_field(self.f_bignum) + self.s.append_field(self.f_testchar) + self.s.append_field(self.f_testushort) + self.s.append_field(self.f_testshort) + self.s.append_field(self.f_testuint) + self.s.append_field(self.f_testint) + self.s.append_field(self.f_testfloat) + self.s.append_field(self.f_testulong) def test_autosize(self): with self.assertRaises(ValueError): @@ -242,3 +262,85 @@ class StructTest(unittest.TestCase): self.assertAlmostEqual(2.8, val[0]) self.assertAlmostEqual(2.79, val[1]) fileobj.seek.assert_called_with(4144, os.SEEK_CUR) + + def test_char_field_set(self): + fileobj = mock.MagicMock(io.BufferedReader) + value = 255 + expected = struct.pack(b">B", value) + self.s.field_set(self.FakeHeader(), fileobj, b"testchar", value) + fileobj.write.assert_called_with(expected) + + with self.assertRaises(struct.error): + self.s.field_set(self.FakeHeader(), fileobj, b"testchar", -1) + + with self.assertRaises(struct.error): + self.s.field_set(self.FakeHeader(), fileobj, b"testchar", 256) + + def test_ushort_field_set(self): + fileobj = mock.MagicMock(io.BufferedReader) + value = 65535 + expected = struct.pack(b">H", value) + self.s.field_set(self.FakeHeader(), fileobj, b"testushort", value) + fileobj.write.assert_called_with(expected) + + with self.assertRaises(struct.error): + self.s.field_set(self.FakeHeader(), fileobj, b"testushort", -1) + + with self.assertRaises(struct.error): + self.s.field_set(self.FakeHeader(), fileobj, b"testushort", 65536) + + def test_short_field_set(self): + fileobj = mock.MagicMock(io.BufferedReader) + value = 32767 + expected = struct.pack(b">h", value) + self.s.field_set(self.FakeHeader(), fileobj, b"testshort", value) + fileobj.write.assert_called_with(expected) + + value = -32768 + expected = struct.pack(b">h", value) + self.s.field_set(self.FakeHeader(), fileobj, b"testshort", value) + fileobj.write.assert_called_with(expected) + + with self.assertRaises(struct.error): + self.s.field_set(self.FakeHeader(), fileobj, b"testshort", -32769) + + with self.assertRaises(struct.error): + self.s.field_set(self.FakeHeader(), fileobj, b"testshort", 32768) + + def test_uint_field_set(self): + fileobj = mock.MagicMock(io.BufferedReader) + value = 4294967295 + expected = struct.pack(b">I", value) + self.s.field_set(self.FakeHeader(), fileobj, b"testuint", value) + fileobj.write.assert_called_with(expected) + + with self.assertRaises(struct.error): + self.s.field_set(self.FakeHeader(), fileobj, b"testuint", -1) + + with self.assertRaises(struct.error): + self.s.field_set(self.FakeHeader(), fileobj, b"testuint", 4294967296) + + def test_int_field_set(self): + fileobj = mock.MagicMock(io.BufferedReader) + value = 2147483647 + expected = struct.pack(b">i", value) + self.s.field_set(self.FakeHeader(), fileobj, b"testint", value) + fileobj.write.assert_called_with(expected) + + value = -2147483648 + expected = struct.pack(b">i", value) + self.s.field_set(self.FakeHeader(), fileobj, b"testint", value) + fileobj.write.assert_called_with(expected) + + with self.assertRaises(struct.error): + self.s.field_set(self.FakeHeader(), fileobj, b"testint", -2147483649) + + with self.assertRaises(struct.error): + self.s.field_set(self.FakeHeader(), fileobj, b"testint", 2147483649) + + def test_float_field_set(self): + fileobj = mock.MagicMock(io.BufferedReader) + value = 3.402823466e38 + expected = struct.pack(b">f", value) + self.s.field_set(self.FakeHeader(), fileobj, b"testfloat", value) + fileobj.write.assert_called_with(expected)