build_psk now creates an armature tree and uses that instead of a flat armature list

io_scene_psk_psa/psk/builder.py

@@ -8,6 +8,8 @@ from psk_psa_py.shared.data import Vector3
 from psk_psa_py.psk.data import Psk
 from .properties import triangle_type_and_bit_flags_to_poly_flags
 from ..shared.helpers import (
+    ObjectNode,
+    ObjectTree,
     PskInputObjects,
     PsxBoneCollection,
     convert_string_to_cp1252_bytes,
@@ -38,11 +40,13 @@ class PskBuildResult(object):
         self.warnings: List[str] = warnings
 
 
-def _get_mesh_export_space_matrix(armature_object: Object | None, export_space: str, root_armature_object: Object | None) -> Matrix:
-    # TODO: this should be a bundle of armature objects. otherwise this creates a scenario where you can have 
-    if armature_object is None or root_armature_object is None:
+def _get_mesh_export_space_matrix(node: ObjectNode | None, export_space: str) -> Matrix:
+    if node is None:
         return Matrix.Identity(4)
     
+    armature_object = node.object
+    root_armature_object = node.root.object
+
     def get_object_space_matrix(obj: Object) -> Matrix:
         translation, rotation, _ = obj.matrix_world.decompose()
         # We neutralize the scale here because the scale is already applied to the mesh objects implicitly.
@@ -86,6 +90,7 @@ def build_psk(context: Context, input_objects: PskInputObjects, options: PskBuil
     assert context.window_manager
 
     armature_objects = list(input_objects.armature_objects)
+    armature_object_tree = ObjectTree(input_objects.armature_objects)    
 
     warnings: list[str] = []
     psk = Psk()
@@ -138,18 +143,19 @@ def build_psk(context: Context, input_objects: PskInputObjects, options: PskBuil
 
     context.window_manager.progress_begin(0, len(input_objects.mesh_dfs_objects))
     coordinate_system_matrix = get_coordinate_system_transform(options.forward_axis, options.up_axis)
-    root_armature_object = next(iter(input_objects.armature_objects), None)
+    root_armature_object = next(iter(armature_object_tree), None)
 
     # Calculate the export spaces for the armature objects.
     # This is used later to transform the mesh object geometry into the export space.
     armature_mesh_export_space_matrices: dict[Object | None, Matrix] = {None: Matrix.Identity(4)}
+
     if options.export_space == 'ARMATURE':
         # For meshes without an armature modifier, we need to set the export space to the first armature object.
-        armature_mesh_export_space_matrices[None] = _get_mesh_export_space_matrix(root_armature_object, options.export_space, root_armature_object)
+        armature_mesh_export_space_matrices[None] = _get_mesh_export_space_matrix(root_armature_object, options.export_space)
     
-    for armature_object in armature_objects:
-        # TODO: also handle the case of multiple roots; dont' just assume we have one!
-        armature_mesh_export_space_matrices[armature_object] = _get_mesh_export_space_matrix(armature_object, options.export_space, root_armature_object)
+    # TODO: also handle the case of multiple roots; dont' just assume we have one!
+    for armature_node in iter(armature_object_tree):
+        armature_mesh_export_space_matrices[armature_node.object] = _get_mesh_export_space_matrix(armature_node, options.export_space)
     
     # Temporarily force the armature into the rest position.
     # The original pose position setting will be restored at the end.

io_scene_psk_psa/shared/helpers.py

@@ -334,41 +334,47 @@ class PsxBoneCollection:
 class ObjectNode:
     def __init__(self, obj: Object):
         self.object = obj
-        self.children: List['ObjectNode'] = []
+        self.parent: ObjectNode | None = None
+        self.children: List[ObjectNode] = []
+    
+    @property
+    def root(self):
+        """
+        Gets the root in the object hierarchy. This can return itself if this node has no parent.
+        """
+        n = self
+        while n.parent is not None:
+            n = n.parent
+        return n
 
 
 class ObjectTree:
-    def __init__(self) -> None:
+    '''
+    A tree of the armature objects based on their hierarchy.
+    '''
+    def __init__(self, objects: Iterable[Object]):
         self.root_nodes: List[ObjectNode] = []
-    
-    @staticmethod
-    def from_objects(objects: Iterable[Object]) -> 'ObjectTree':
-        '''
-        Make a tree of the armature objects based on their hierarchy.
-        '''
-        tree = ObjectTree()
         object_node_map: Dict[Object, ObjectNode] = {x: ObjectNode(x) for x in objects}
         
         for obj, object_node in object_node_map.items():
             if obj.parent in object_node_map:
                 parent_node = object_node_map[obj.parent]
+                object_node.parent = parent_node
                 parent_node.children.append(object_node)
             else:
-                tree.root_nodes.append(object_node)
-
-        return tree
+                self.root_nodes.append(object_node)
 
     def __iter__(self):
         """
         An depth-first iterator over the armature tree.
         """
-        node_stack = self.root_nodes
+        node_stack = [] + self.root_nodes
         while node_stack:
             node = node_stack.pop()
             yield node
             node_stack = node.children + node_stack
     
-    def objects_dfs(self):
+    def objects_iterator(self):
         for node in self:
             yield node.object
     
@@ -401,7 +407,7 @@ def create_psx_bones(
     if bone_collection_indices is None:
         bone_collection_indices = []
 
-    armature_tree = ObjectTree.from_objects(armature_objects)
+    armature_tree = ObjectTree(armature_objects)
 
     # Check that there is only one root bone. If there are multiple armature objects, the export space must be WORLD.
     if len(armature_tree.root_nodes) >= 2 and export_space != 'WORLD':
@@ -655,7 +661,7 @@ def _get_psk_input_objects(mesh_dfs_objects: Iterable[DfsObject]) -> PskInputObj
     # Get the armature objects used on all the meshes being exported.
     armature_objects = get_armatures_for_mesh_objects(map(lambda x: x.obj, mesh_dfs_objects))
     # Sort them in hierarchy order.
-    input_objects.armature_objects = list(ObjectTree.from_objects(armature_objects).objects_dfs())
+    input_objects.armature_objects = list(ObjectTree(armature_objects).objects_iterator())
     return input_objects