From 3001501006cab02741abcab1553e3898fa6cbfa7 Mon Sep 17 00:00:00 2001
From: Colin Basnett <colin@spexigeo.com>
Date: Thu, 12 Dec 2019 15:07:22 -0800
Subject: [PATCH] Action selection now works in the PSA exporter, and actions
 that likely correspond to the selected armature are deselected by default
 (will save a bit of time!)

---
 src/__init__.py     |  5 +++++
 src/psa/builder.py  |  9 +++++++--
 src/psa/operator.py | 47 +++++++++++++++++++++++++++++++++++++--------
 3 files changed, 51 insertions(+), 10 deletions(-)

diff --git a/src/__init__.py b/src/__init__.py
index 65b10a9..b515278 100644
--- a/src/__init__.py
+++ b/src/__init__.py
@@ -42,12 +42,15 @@ classes = [
     psa_operator.ActionListItem
 ]
 
+
 def psk_menu_func(self, context):
     self.layout.operator(psk_operator.PskExportOperator.bl_idname, text ='Unreal PSK (.psk)')
 
+
 def psa_menu_func(self, context):
     self.layout.operator(psa_operator.PsaExportOperator.bl_idname, text='Unreal PSA (.psa)')
 
+
 def register():
     from bpy.utils import register_class
     for cls in classes:
@@ -57,6 +60,7 @@ def register():
     bpy.types.Scene.psa_action_list = CollectionProperty(type=psa_operator.ActionListItem)
     bpy.types.Scene.psa_action_list_index = IntProperty(name='index for list??', default=0)
 
+
 def unregister():
     del bpy.types.Scene.psa_action_list_index
     del bpy.types.Scene.psa_action_list
@@ -66,5 +70,6 @@ def unregister():
     for cls in reversed(classes):
         unregister_class(cls)
 
+
 if __name__ == '__main__':
     register()
diff --git a/src/psa/builder.py b/src/psa/builder.py
index 4867acc..5a62a3b 100644
--- a/src/psa/builder.py
+++ b/src/psa/builder.py
@@ -3,13 +3,18 @@ import mathutils
 from .data import *
 
 
+class PsaBuilderOptions(object):
+    def __init__(self):
+        self.actions = []
+
+
 # https://git.cth451.me/cth451/blender-addons/blob/master/io_export_unreal_psk_psa.py
 class PsaBuilder(object):
     def __init__(self):
         # TODO: add options in here (selected anims, eg.)
         pass
 
-    def build(self, context) -> Psa:
+    def build(self, context, options) -> Psa:
         object = context.view_layer.objects.active
 
         if object.type != 'ARMATURE':
@@ -70,7 +75,7 @@ class PsaBuilder(object):
         print('---- ACTIONS ----')
         frame_start_index = 0
 
-        for action in bpy.data.actions:
+        for action in options.actions:
             if len(action.fcurves) == 0:
                 continue
 
diff --git a/src/psa/operator.py b/src/psa/operator.py
index 2ed7fc7..4dd7a96 100644
--- a/src/psa/operator.py
+++ b/src/psa/operator.py
@@ -1,7 +1,7 @@
 from bpy.types import Operator, Action, UIList, PropertyGroup
 from bpy_extras.io_utils import ExportHelper
 from bpy.props import StringProperty, BoolProperty, FloatProperty, CollectionProperty, PointerProperty
-from .builder import PsaBuilder
+from .builder import PsaBuilder, PsaBuilderOptions
 from .exporter import PsaExporter
 import bpy
 import re
@@ -20,7 +20,7 @@ class PSA_UL_ActionList(UIList):
     def draw_item(self, context, layout, data, item, icon, active_data, active_propname, index):
         layout.alignment = 'LEFT'
         layout.prop(item, 'is_selected', icon_only=True)
-        layout.label(text=item.action.name, icon='ACTION')
+        layout.label(text=item.action.name)
 
     def filter_items(self, context, data, property):
         # TODO: returns two lists, apparently
@@ -44,30 +44,61 @@ class PsaExportOperator(Operator, ExportHelper):
         maxlen=1024,
         default='')
 
+    def __init__(self):
+        self.armature = None
+
     def draw(self, context):
         layout = self.layout
         scene = context.scene
-        row = layout.row()
-        row.label(text='Actions')
-        row = layout.row()
+        box = layout.box()
+        box.label(text='Actions', icon='ACTION')
+        row = box.row()
         row.template_list('PSA_UL_ActionList', 'asd', scene, 'psa_action_list', scene, 'psa_action_list_index', rows=len(context.scene.psa_action_list))
 
+    def is_action_for_armature(self, action):
+        bone_names = [x.name for x in self.armature.data.bones]
+        print(bone_names)
+        for fcurve in action.fcurves:
+            match = re.match('pose\.bones\["(.+)"\].\w+', fcurve.data_path)
+            if not match:
+                continue
+            bone_name = match.group(1)
+            if bone_name not in bone_names:
+                return False
+        return True
+
     def invoke(self, context, event):
         if context.view_layer.objects.active.type != 'ARMATURE':
             self.report({'ERROR_INVALID_CONTEXT'}, 'The selected object must be an armature.')
             return {'CANCELLED'}
+
+        self.armature = context.view_layer.objects.active
+
         context.scene.psa_action_list.clear()
         for action in bpy.data.actions:
             item = context.scene.psa_action_list.add()
             item.action = action
-            # TODO: add
-            item.is_selected = True
+            if self.is_action_for_armature(action):
+                item.is_selected = True
+
+        if len(context.scene.psa_action_list) == 0:
+            self.report({'ERROR_INVALID_CONTEXT'}, 'There are no actions to export.')
+            return {'CANCELLED'}
+
         context.window_manager.fileselect_add(self)
         return {'RUNNING_MODAL'}
 
     def execute(self, context):
+        actions = [x.action for x in context.scene.psa_action_list if x.is_selected]
+
+        if len(actions) == 0:
+            self.report({'ERROR_INVALID_CONTEXT'}, 'No actions were selected for export.')
+            return {'CANCELLED'}
+
+        options = PsaBuilderOptions()
+        options.actions = actions
         builder = PsaBuilder()
-        psk = builder.build(context)
+        psk = builder.build(context, options)
         exporter = PsaExporter(psk)
         exporter.export(self.filepath)
         return {'FINISHED'}