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)