Merge initial a3d export support

Initial a3d export support

Author: Pyogenics

io_scene_a3d/A3D.py

@@ -20,7 +20,9 @@
 SOFTWARE.
 '''
 
-from .IOTools import unpackStream, readNullTerminatedString, calculatePadding
+from io import BytesIO
+
+from .IOTools import unpackStream, packStream, readNullTerminatedString, calculatePadding
 from . import A3DObjects
 
 '''
@@ -65,13 +67,33 @@ def read(self, stream):
             self.readRootBlock2(stream)
         elif self.version == 3:
             self.readRootBlock3(stream)
+        else:
+            raise RuntimeError(f"Unknown A3D version: {self.version}")
+
+    def write(self, stream, version=2):
+        # Write header data
+        stream.write(A3D_SIGNATURE)
+        packStream("<2H", stream, version, 0)
+        
+        # Write root block
+        if version == 1:
+            self.writeRootBlock1(stream)
+        elif version == 2:
+            self.writeRootBlock2(stream)
+        elif version == 3:
+            self.writeRootBlock3(stream)
+        else:
+            raise RuntimeError(f"Unknown A3D version: {version} whilst writing A3D")
 
     '''
     Root data blocks
     '''
     def readRootBlock1(self, stream):
         raise RuntimeError("Version 1 files are not supported yet")
 
+    def writeRootBlock1(self, stream):
+        raise RuntimeError("Version 1 files are not supported yet")
+
     def readRootBlock2(self, stream):
         # Verify signature
         signature, _ = unpackStream("<2I", stream)
@@ -85,6 +107,21 @@ def readRootBlock2(self, stream):
         self.readTransformBlock2(stream)
         self.readObjectBlock2(stream)
 
+    def writeRootBlock2(self, stream):
+        buffer = BytesIO()
+
+        # Write data to the buffer
+        print("Writing root block")
+        self.writeMaterialBlock2(buffer)
+        self.writeMeshBlock2(buffer)
+        self.writeTransformBlock2(buffer)
+        self.writeObjectBlock2(buffer)
+
+        # Write buffer to stream
+        packStream("<2I", stream, A3D_ROOTBLOCK_SIGNATURE, buffer.tell())
+        buffer.seek(0, 0)
+        stream.write(buffer.read())
+
     def readRootBlock3(self, stream):
         # Verify signature
         signature, length = unpackStream("<2I", stream)
@@ -101,6 +138,21 @@ def readRootBlock3(self, stream):
         padding = calculatePadding(length)
         stream.read(padding)
 
+    def writeRootBlock3(self, stream):
+        buffer = BytesIO()
+
+        # Write data to the buffer
+        print("Writing root block")
+        self.writeMaterialBlock3(buffer)
+        self.writeMeshBlock3(buffer)
+        self.writeTransformBlock3(buffer)
+        self.writeObjectBlock3(buffer)
+
+        # Write buffer to stream
+        packStream("<2I", stream, A3D_ROOTBLOCK_SIGNATURE, buffer.tell())
+        buffer.seek(0, 0)
+        stream.write(buffer.read())
+
     '''
     Material data blocks
     '''
@@ -117,6 +169,20 @@ def readMaterialBlock2(self, stream):
             material.read2(stream)
             self.materials.append(material)
 
+    def writeMaterialBlock2(self, stream):
+        buffer = BytesIO()
+
+        # Write data to the buffer
+        print("Writing material block")
+        packStream("<I", buffer, len(self.materials))
+        for material in self.materials:
+            material.write2(buffer)
+
+        # Write buffer to stream
+        packStream("<2I", stream, A3D_MATERIALBLOCK_SIGNATURE, buffer.tell())
+        buffer.seek(0, 0)
+        stream.write(buffer.read()) 
+    
     def readMaterialBlock3(self, stream):
         # Verify signature
         signature, length, materialCount = unpackStream("<3I", stream)
@@ -134,6 +200,24 @@ def readMaterialBlock3(self, stream):
         padding = calculatePadding(length)
         stream.read(padding)
 
+    def writeMaterialBlock3(self, stream):
+        buffer = BytesIO()
+
+        # Write data to the buffer
+        print("Writing material block")
+        packStream("<I", buffer, len(self.materials))
+        for material in self.materials:
+            material.write3(buffer)
+
+        # Write buffer to stream
+        packStream("<2I", stream, A3D_MATERIALBLOCK_SIGNATURE, buffer.tell())
+        buffer.seek(0, 0)
+        stream.write(buffer.read())
+
+        # Padding
+        paddingSize = calculatePadding(buffer.tell())
+        stream.write(b"\x00" * paddingSize)
+
     '''
     Mesh data blocks
     '''
@@ -150,6 +234,20 @@ def readMeshBlock2(self, stream):
             mesh.read2(stream)
             self.meshes.append(mesh)
 
+    def writeMeshBlock2(self, stream):
+        buffer = BytesIO()
+
+        # Write data to the buffer
+        print("Writing mesh block")
+        packStream("<I", buffer, len(self.meshes))
+        for mesh in self.meshes:
+            mesh.write2(buffer)
+
+        # Write buffer to stream
+        packStream("<2I", stream, A3D_MESHBLOCK_SIGNATURE, buffer.tell())
+        buffer.seek(0, 0)
+        stream.write(buffer.read())
+
     def readMeshBlock3(self, stream):
         # Verify signature
         signature, length, meshCount = unpackStream("<3I", stream)
@@ -167,6 +265,24 @@ def readMeshBlock3(self, stream):
         padding = calculatePadding(length)
         stream.read(padding)
 
+    def writeMeshBlock3(self, stream):
+        buffer = BytesIO()
+
+        # Write data to the buffer
+        print("Writing mesh block")
+        packStream("<I", buffer, len(self.meshes))
+        for mesh in self.meshes:
+            mesh.write3(buffer)
+
+        # Write buffer to stream
+        packStream("<2I", stream, A3D_MESHBLOCK_SIGNATURE, buffer.tell())
+        buffer.seek(0, 0)
+        stream.write(buffer.read())
+
+        # Padding
+        paddingSize = calculatePadding(buffer.tell())
+        stream.write(b"\x00" * paddingSize)
+
     '''
     Transform data blocks
     '''
@@ -187,6 +303,22 @@ def readTransformBlock2(self, stream):
             parentID, = unpackStream("<i", stream)
             self.transformParentIDs.append(parentID)
 
+    def writeTransformBlock2(self, stream):
+        buffer = BytesIO()
+
+        # Write data to the buffer
+        print("Writing transform block")
+        packStream("<I", buffer, len(self.transforms))
+        for transform in self.transforms:
+            transform.write2(buffer)
+        for parentID in self.transformParentIDs:
+            packStream("<i", buffer, parentID)
+
+        # Write buffer to stream
+        packStream("<2I", stream, A3D_TRANSFORMBLOCK_SIGNATURE, buffer.tell())
+        buffer.seek(0, 0)
+        stream.write(buffer.read())
+
     def readTransformBlock3(self, stream):
         # Verify signature
         signature, length, transformCount = unpackStream("<3I", stream)
@@ -195,7 +327,6 @@ def readTransformBlock3(self, stream):
 
         # Read data
         print(f"Reading transform block with {transformCount} transforms and length {length}")
-        transforms = []
         for _ in range(transformCount):
             transform = A3DObjects.A3DTransform()
             transform.read3(stream)
@@ -209,6 +340,26 @@ def readTransformBlock3(self, stream):
         padding = calculatePadding(length)
         stream.read(padding)
 
+    def writeTransformBlock3(self, stream):
+        buffer = BytesIO()
+
+        # Write data to the buffer
+        print("Writing transform block")
+        packStream("<I", buffer, len(self.transforms))
+        for transform in self.transforms:
+            transform.write3(buffer)
+        for parentID in self.transformParentIDs:
+            packStream("<i", buffer, parentID)
+
+        # Write buffer to stream
+        packStream("<2I", stream, A3D_TRANSFORMBLOCK_SIGNATURE, buffer.tell())
+        buffer.seek(0, 0)
+        stream.write(buffer.read())
+
+        # Padding
+        paddingSize = calculatePadding(buffer.tell())
+        stream.write(b"\x00" * paddingSize)
+
     '''
     Object data blocks
     '''
@@ -225,6 +376,20 @@ def readObjectBlock2(self, stream):
             objec.read2(stream)
             self.objects.append(objec)
 
+    def writeObjectBlock2(self, stream):
+        buffer = BytesIO()
+
+        # Write data to the buffer
+        print("Writing object block")
+        packStream("<I", buffer, len(self.objects))
+        for objec in self.objects:
+            objec.write2(buffer)
+
+        # Write buffer to stream
+        packStream("<2I", stream, A3D_OBJECTBLOCK_SIGNATURE, buffer.tell())
+        buffer.seek(0, 0)
+        stream.write(buffer.read())
+
     def readObjectBlock3(self, stream):
         # Verify signature
         signature, length, objectCount = unpackStream("<3I", stream)
@@ -240,4 +405,22 @@ def readObjectBlock3(self, stream):
 
         # Padding
         padding = calculatePadding(length)
-        stream.read(padding)
+        stream.read(padding)
+
+    def writeObjectBlock3(self, stream):
+        buffer = BytesIO()
+
+        # Write data to the buffer
+        print("Writing object block")
+        packStream("<I", buffer, len(self.objects))
+        for objec in self.objects:
+            objec.write3(buffer)
+
+        # Write buffer to stream
+        packStream("<2I", stream, A3D_OBJECTBLOCK_SIGNATURE, buffer.tell())
+        buffer.seek(0, 0)
+        stream.write(buffer.read())
+
+        # Padding
+        paddingSize = calculatePadding(buffer.tell())
+        stream.write(b"\x00" * paddingSize)