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

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

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

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

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)