Added "Export Space" option to allow exporting in armature space

Armature space exporting means that the user doesn't need to manually
move the armature back to origin before exporting, since it was simply
using world-space exporting before.

Also moved all DFS functions into their own file & cleaned up the
usage interface. Added `is_visible` and `is_selected` properties
for better ergonomics.

io_scene_psk_psa/psk/builder.py

@@ -1,19 +1,20 @@
 import typing
-from typing import Optional, Set
+from typing import Optional
 
 import bmesh
 import numpy as np
-from bpy.types import Material, Collection, Context, LayerCollection, ViewLayer
+from bpy.types import Material, Collection, Context, Mesh
 from mathutils import Matrix
 
 from .data import *
 from .properties import triangle_type_and_bit_flags_to_poly_flags
+from ..shared.dfs import dfs_collection_objects, dfs_view_layer_objects, DfsObject
 from ..shared.helpers import *
 
 
 class PskInputObjects(object):
     def __init__(self):
-        self.mesh_objects: List[Tuple[Object, List[Object], Matrix]] = []
+        self.mesh_objects: List[DfsObject] = []
         self.armature_object: Optional[Object] = None
 
 
@@ -25,26 +26,17 @@ class PskBuildOptions(object):
         self.materials: List[Material] = []
         self.should_enforce_bone_name_restrictions = False
         self.scale = 1.0
+        self.export_space = 'WORLD'
 
 
-def get_mesh_objects_for_collection(collection: Collection, should_exclude_hidden_meshes: bool = True) -> Iterable[Tuple[Object, List[Object], Matrix]]:
+def get_mesh_objects_for_collection(collection: Collection) -> Iterable[DfsObject]:
-    for obj, instance_objects, matrix in dfs_collection_objects(collection):
+    return filter(lambda x: x.obj.type == 'MESH', dfs_collection_objects(collection))
-        if obj.type != 'MESH':
-            continue
-        if should_exclude_hidden_meshes:
-            if instance_objects:
-                if not instance_objects[-1].visible_get():
-                    continue
-            elif not obj.visible_get():
-                continue
-        yield (obj, instance_objects, matrix)
 
 
-def get_mesh_objects_for_context(context: Context) -> Iterable[Tuple[Object, List[Object], Matrix]]:
+def get_mesh_objects_for_context(context: Context) -> Iterable[DfsObject]:
-    for (obj, instance_objects, matrix) in dfs_view_layer_objects(context.view_layer):
+    for dfs_object in dfs_view_layer_objects(context.view_layer):
-        is_selected = obj.select_get() or any(x.select_get() for x in instance_objects)
+        if dfs_object.obj.type == 'MESH' and dfs_object.is_selected:
-        if obj.type == 'MESH' and is_selected:
+            yield dfs_object
-            yield (obj, instance_objects, matrix)
 
 
 def get_armature_for_mesh_objects(mesh_objects: Iterable[Object]) -> Optional[Object]:
@@ -69,17 +61,19 @@ def get_armature_for_mesh_objects(mesh_objects: Iterable[Object]) -> Optional[Ob
         return None
 
 
-def _get_psk_input_objects(mesh_objects: List[Tuple[Object, List[Object], Matrix]]) -> PskInputObjects:
+def _get_psk_input_objects(mesh_objects: Iterable[DfsObject]) -> PskInputObjects:
+    mesh_objects = list(mesh_objects)
     if len(mesh_objects) == 0:
         raise RuntimeError('At least one mesh must be selected')
 
-    for mesh_object, _, _ in mesh_objects:
+    for dfs_object in mesh_objects:
-        if len(mesh_object.data.materials) == 0:
+        mesh_data = cast(Mesh, dfs_object.obj.data)
-            raise RuntimeError(f'Mesh "{mesh_object.name}" must have at least one material')
+        if len(mesh_data.materials) == 0:
+            raise RuntimeError(f'Mesh "{dfs_object.obj.name}" must have at least one material')
 
     input_objects = PskInputObjects()
     input_objects.mesh_objects = mesh_objects
-    input_objects.armature_object = get_armature_for_mesh_objects([x[0] for x in mesh_objects])
+    input_objects.armature_object = get_armature_for_mesh_objects([x.obj for x in mesh_objects])
 
     return input_objects
 
@@ -90,7 +84,9 @@ def get_psk_input_objects_for_context(context: Context) -> PskInputObjects:
 
 
 def get_psk_input_objects_for_collection(collection: Collection, should_exclude_hidden_meshes: bool = True) -> PskInputObjects:
-    mesh_objects = list(get_mesh_objects_for_collection(collection, should_exclude_hidden_meshes))
+    mesh_objects = get_mesh_objects_for_collection(collection)
+    if should_exclude_hidden_meshes:
+        mesh_objects = filter(lambda x: x.is_visible, mesh_objects)
     return _get_psk_input_objects(mesh_objects)
 
 
@@ -107,6 +103,16 @@ def build_psk(context, input_objects: PskInputObjects, options: PskBuildOptions)
     psk = Psk()
     bones = []
 
+    def get_export_space_matrix():
+        match options.export_space:
+            case 'WORLD':
+                return Matrix.Identity(4)
+            case 'ARMATURE':
+                return armature_object.matrix_world.inverted()
+            case _:
+                raise ValueError(f'Invalid export space: {options.export_space}')
+
+    export_space_matrix = get_export_space_matrix()  # TODO: maybe neutralize the scale here?
     scale_matrix = Matrix.Scale(options.scale, 4)
 
     if armature_object is None or len(armature_object.data.bones) == 0:
@@ -153,7 +159,16 @@ def build_psk(context, input_objects: PskInputObjects, options: PskBuildOptions)
                 parent_tail = inverse_parent_rotation @ bone.parent.tail
                 location = (parent_tail - parent_head) + bone.head
             else:
-                armature_local_matrix = armature_object.matrix_local
+                def get_armature_local_matrix():
+                    match options.export_space:
+                        case 'WORLD':
+                            return armature_object.matrix_world
+                        case 'ARMATURE':
+                            return Matrix.Identity(4)
+                        case _:
+                            raise ValueError(f'Invalid export space: {options.export_space}')
+
+                armature_local_matrix = get_armature_local_matrix()
                 location = armature_local_matrix @ bone.head
                 bone_rotation = bone.matrix.to_quaternion().conjugated()
                 local_rotation = armature_local_matrix.to_3x3().to_quaternion().conjugated()
@@ -162,6 +177,14 @@ def build_psk(context, input_objects: PskInputObjects, options: PskBuildOptions)
 
             location = scale_matrix @ location
 
+            # If the armature object has been scaled, we need to scale the bone's location to match.
+            _, _, armature_object_scale = armature_object.matrix_world.decompose()
+            location.x *= armature_object_scale.x
+            location.y *= armature_object_scale.y
+            location.z *= armature_object_scale.z
+
+            print(bone.name, location)
+
             psk_bone.location.x = location.x
             psk_bone.location.y = location.y
             psk_bone.location.z = location.z
@@ -171,12 +194,6 @@ def build_psk(context, input_objects: PskInputObjects, options: PskBuildOptions)
             psk_bone.rotation.y = rotation.y
             psk_bone.rotation.z = rotation.z
 
-            # If the armature object has been scaled, we need to scale the bone's location to match.
-            _, _, armature_object_scale = armature_object.matrix_world.decompose()
-            psk_bone.location.x *= armature_object_scale.x
-            psk_bone.location.y *= armature_object_scale.y
-            psk_bone.location.z *= armature_object_scale.z
-
             psk.bones.append(psk_bone)
 
     # MATERIALS
@@ -197,7 +214,7 @@ def build_psk(context, input_objects: PskInputObjects, options: PskBuildOptions)
 
     for object_index, input_mesh_object in enumerate(input_objects.mesh_objects):
 
-        obj, instance_objects, matrix_world = input_mesh_object
+        obj, instance_objects, matrix_world = input_mesh_object.obj, input_mesh_object.instance_objects, input_mesh_object.matrix_world
 
         should_flip_normals = False
 
@@ -252,7 +269,7 @@ def build_psk(context, input_objects: PskInputObjects, options: PskBuildOptions)
                 raise ValueError(f'Invalid object evaluation state: {options.object_eval_state}')
 
         vertex_offset = len(psk.points)
-        matrix_world = scale_matrix @ mesh_object.matrix_world
+        matrix_world = scale_matrix @ export_space_matrix @ mesh_object.matrix_world
 
         # VERTICES
         for vertex in mesh_data.vertices:
@@ -383,75 +400,3 @@ def build_psk(context, input_objects: PskInputObjects, options: PskBuildOptions)
     result.psk = psk
 
     return result
-
-
-def dfs_collection_objects_recursive(collection: Collection, visited: Optional[Set[Object]]=None) -> Iterable[Tuple[Object, List[Object], Matrix]]:
-    if visited is None:
-        visited = set()
-    yield from dfs_collection_objects(collection, visited=visited)
-    for child in collection.children:
-        yield from dfs_collection_objects_recursive(child, visited)
-
-
-# Construct a list of objects in hierarchy order from `collection.objects`, only keeping those that are in the
-# collection.
-def dfs_object_children(obj: Object, collection: Collection):
-    yield obj
-    for child in obj.children:
-        if child in collection.objects:
-            yield from dfs_object_children(child, collection)
-
-
-def dfs_objects_in_collection(collection: Collection):
-    # Return only the top-level objects in the collection.
-    objects_hierarchy = []
-    for obj in collection.objects:
-        if obj.parent is None or obj.parent not in set(collection.objects):
-            objects_hierarchy.append(obj)
-    for obj in collection.objects:
-        yield from dfs_object_children(obj, collection)
-
-
-def dfs_collection_objects(
-        collection: Collection,
-        instance_objects: Optional[List[Object]] = None,
-        matrix_world: Matrix = Matrix.Identity(4),
-        visited: Optional[Set[Object]]=None
-) -> Iterable[Tuple[Object, List[Object], Matrix]]:
-    # We want to also yield the top-level instance object so that callers can inspect the selection status etc.
-    if visited is None:
-        visited = set()
-
-    if instance_objects is None:
-        instance_objects = list()
-
-    for child in collection.children:
-        yield from dfs_collection_objects(child, instance_objects, matrix_world.copy(), visited)
-
-    for obj in dfs_objects_in_collection(collection):
-        visited_pair = (obj, instance_objects[-1] if instance_objects else None)
-        if visited_pair in visited:
-            continue
-        # If this an instance, we need to recurse into it.
-        if obj.instance_collection is not None:
-            # Calculate the instance transform.
-            instance_offset_matrix = Matrix.Translation(-obj.instance_collection.instance_offset)
-            # Recurse into the instance collection.
-            yield from dfs_collection_objects(obj.instance_collection,
-                                     instance_objects + [obj],
-                                     matrix_world @ (obj.matrix_world @ instance_offset_matrix),
-                                     visited)
-        else:
-            # Object is not an instance, yield it.
-            yield (obj, instance_objects, matrix_world @ obj.matrix_world)
-            visited.add(visited_pair)
-
-
-def dfs_view_layer_objects(view_layer: ViewLayer) -> Iterable[Tuple[Object, List[Object], Matrix]]:
-    def layer_collection_objects_recursive(layer_collection: LayerCollection):
-        for child in layer_collection.children:
-            yield from layer_collection_objects_recursive(child)
-        # Iterate only the top-level objects in this collection first.
-        yield from dfs_collection_objects(layer_collection.collection)
-
-    yield from layer_collection_objects_recursive(view_layer.layer_collection)

io_scene_psk_psa/psk/export/operators.py

@@ -5,7 +5,7 @@ from bpy.props import StringProperty, BoolProperty, EnumProperty, FloatProperty
 from bpy.types import Operator, Context, Object
 from bpy_extras.io_utils import ExportHelper
 
-from .properties import object_eval_state_items
+from .properties import object_eval_state_items, export_space_items
 from ..builder import build_psk, PskBuildOptions, get_psk_input_objects_for_context, \
     get_psk_input_objects_for_collection
 from ..writer import write_psk
@@ -115,6 +115,13 @@ class PSK_OT_export_collection(Operator, ExportHelper):
         min=0.0001,
         soft_max=100.0
     )
+    export_space: EnumProperty(
+        name='Export Space',
+        description='Space to export the mesh in',
+        items=export_space_items,
+        default='WORLD'
+    )
+
 
     def execute(self, context):
         collection = bpy.data.collections.get(self.collection)
@@ -128,9 +135,10 @@ class PSK_OT_export_collection(Operator, ExportHelper):
         options = PskBuildOptions()
         options.bone_filter_mode = 'ALL'
         options.object_eval_state = self.object_eval_state
-        options.materials = get_materials_for_mesh_objects([x[0] for x in input_objects.mesh_objects])
+        options.materials = get_materials_for_mesh_objects([x.obj for x in input_objects.mesh_objects])
         options.should_enforce_bone_name_restrictions = self.should_enforce_bone_name_restrictions
         options.scale = self.scale
+        options.export_space = self.export_space
 
         try:
             result = build_psk(context, input_objects, options)
@@ -155,6 +163,7 @@ class PSK_OT_export_collection(Operator, ExportHelper):
         flow.use_property_decorate = False
 
         flow.prop(self, 'scale')
+        flow.prop(self, 'export_space')
 
         # MESH
         mesh_header, mesh_panel = layout.panel('Mesh', default_closed=False)
@@ -271,6 +280,7 @@ class PSK_OT_export(Operator, ExportHelper):
         options.materials = [m.material for m in pg.material_list]
         options.should_enforce_bone_name_restrictions = pg.should_enforce_bone_name_restrictions
         options.scale = pg.scale
+        options.export_space = pg.export_space
         
         try:
             result = build_psk(context, input_objects, options)

io_scene_psk_psa/psk/export/properties.py

@@ -11,6 +11,11 @@ object_eval_state_items = (
     ('ORIGINAL', 'Original', 'Use data from original object with no modifiers applied'),
 )
 
+export_space_items = [
+    ('WORLD', 'World', 'Export in world space'),
+    ('ARMATURE', 'Armature', 'Export in armature space'),
+]
+
 class PSK_PG_material_list_item(PropertyGroup):
     material: PointerProperty(type=Material)
     index: IntProperty()
@@ -49,6 +54,13 @@ class PSK_PG_export(PropertyGroup):
         min=0.0001,
         soft_max=100.0
     )
+    export_space: EnumProperty(
+        name='Export Space',
+        options=empty_set,
+        description='Space to export the mesh in',
+        items=export_space_items,
+        default='WORLD'
+    )
 
 
 classes = (

io_scene_psk_psa/psk/import_/operators.py

@@ -152,6 +152,7 @@ class PSK_OT_import(Operator, ImportHelper):
         col.use_property_split = True
         col.use_property_decorate = False
         col.prop(self, 'scale')
+        col.prop(self, 'export_space')
 
         mesh_header, mesh_panel = layout.panel('mesh_panel_id', default_closed=False)
         mesh_header.prop(self, 'should_import_mesh')

io_scene_psk_psa/shared/dfs.py

@@ -0,0 +1,154 @@
+'''
+Depth-first object iterator functions for Blender collections and view layers.
+
+These functions are used to iterate over objects in a collection or view layer in a depth-first manner, including
+instances. This is useful for exporters that need to traverse the object hierarchy in a predictable order.
+'''
+
+from typing import Optional, Set, Iterable, List
+
+from bpy.types import Collection, Object, ViewLayer, LayerCollection
+from mathutils import Matrix
+
+
+class DfsObject:
+    '''
+    Represents an object in a depth-first search.
+    '''
+    def __init__(self, obj: Object, instance_objects: List[Object], matrix_world: Matrix):
+        self.obj = obj
+        self.instance_objects = instance_objects
+        self.matrix_world = matrix_world
+
+    @property
+    def is_visible(self) -> bool:
+        '''
+        Check if the object is visible.
+        @return: True if the object is visible, False otherwise.
+        '''
+        if self.instance_objects:
+            return self.instance_objects[-1].visible_get()
+        return self.obj.visible_get()
+
+    @property
+    def is_selected(self) -> bool:
+        '''
+        Check if the object is selected.
+        @return: True if the object is selected, False otherwise.
+        '''
+        if self.instance_objects:
+            return self.instance_objects[-1].select_get()
+        return self.obj.select_get()
+
+
+
+def _dfs_object_children(obj: Object, collection: Collection) -> Iterable[Object]:
+    '''
+    Construct a list of objects in hierarchy order from `collection.objects`, only keeping those that are in the
+    collection.
+    @param obj: The object to start the search from.
+    @param collection:  The collection to search in.
+    @return: An iterable of objects in hierarchy order.
+    '''
+    yield obj
+    for child in obj.children:
+        if child.name in collection.objects:
+            yield from _dfs_object_children(child, collection)
+
+
+def dfs_objects_in_collection(collection: Collection) -> Iterable[Object]:
+    '''
+    Returns a depth-first iterator over all objects in a collection, only keeping those that are directly in the
+    collection.
+    @param collection: The collection to search in.
+    @return: An iterable of objects in hierarchy order.
+    '''
+    objects_hierarchy = []
+    for obj in collection.objects:
+        if obj.parent is None or obj.parent not in set(collection.objects):
+            objects_hierarchy.append(obj)
+    for obj in objects_hierarchy:
+        yield from _dfs_object_children(obj, collection)
+
+
+def dfs_collection_objects(collection: Collection, visible_only: bool = False) -> Iterable[DfsObject]:
+    '''
+    Depth-first search of objects in a collection, including recursing into instances.
+    @param collection: The collection to search in.
+    @return: An iterable of tuples containing the object, the instance objects, and the world matrix.
+    '''
+    yield from _dfs_collection_objects_recursive(collection)
+
+
+def _dfs_collection_objects_recursive(
+        collection: Collection,
+        instance_objects: Optional[List[Object]] = None,
+        matrix_world: Matrix = Matrix.Identity(4),
+        visited: Optional[Set[Object]]=None
+) -> Iterable[DfsObject]:
+    '''
+    Depth-first search of objects in a collection, including recursing into instances.
+    This is a recursive function.
+    @param collection: The collection to search in.
+    @param instance_objects: The running hierarchy of instance objects.
+    @param matrix_world: The world matrix of the current object.
+    @param visited: A set of visited object-instance pairs.
+    @return: An iterable of tuples containing the object, the instance objects, and the world matrix.
+    '''
+
+    # We want to also yield the top-level instance object so that callers can inspect the selection status etc.
+    if visited is None:
+        visited = set()
+
+    if instance_objects is None:
+        instance_objects = list()
+
+    # First, yield all objects in child collections.
+    for child in collection.children:
+        yield from _dfs_collection_objects_recursive(child, instance_objects, matrix_world.copy(), visited)
+
+    # Then, evaluate all objects in this collection.
+    for obj in dfs_objects_in_collection(collection):
+        visited_pair = (obj, instance_objects[-1] if instance_objects else None)
+        if visited_pair in visited:
+            continue
+        # If this an instance, we need to recurse into it.
+        if obj.instance_collection is not None:
+            # Calculate the instance transform.
+            instance_offset_matrix = Matrix.Translation(-obj.instance_collection.instance_offset)
+            # Recurse into the instance collection.
+            yield from _dfs_collection_objects_recursive(obj.instance_collection,
+                                                         instance_objects + [obj],
+                                                         matrix_world @ (obj.matrix_world @ instance_offset_matrix),
+                                                         visited)
+        else:
+            # Object is not an instance, yield it.
+            yield DfsObject(obj, instance_objects, matrix_world @ obj.matrix_world)
+            visited.add(visited_pair)
+
+
+def dfs_view_layer_objects(view_layer: ViewLayer) -> Iterable[DfsObject]:
+    '''
+    Depth-first iterator over all objects in a view layer, including recursing into instances.
+    @param view_layer: The view layer to inspect.
+    @return: An iterable of tuples containing the object, the instance objects, and the world matrix.
+    '''
+    def layer_collection_objects_recursive(layer_collection: LayerCollection):
+        for child in layer_collection.children:
+            yield from layer_collection_objects_recursive(child)
+        # Iterate only the top-level objects in this collection first.
+        yield from _dfs_collection_objects_recursive(layer_collection.collection)
+
+    yield from layer_collection_objects_recursive(view_layer.layer_collection)
+
+
+def _is_dfs_object_visible(obj: Object, instance_objects: List[Object]) -> bool:
+    '''
+    Check if a DFS object is visible.
+    @param obj: The object.
+    @param instance_objects: The instance objects.
+    @return: True if the object is visible, False otherwise.
+    '''
+    if instance_objects:
+        return instance_objects[-1].visible_get()
+    return obj.visible_get()