From 9a530cf68a79ffeae279f8f513aacb97252966b8 Mon Sep 17 00:00:00 2001
From: Jack Yao <151671689+jack-yao91@users.noreply.github.com>
Date: Tue, 30 Jul 2024 18:05:58 -0500
Subject: [PATCH] UE2Rigify - Custom Rig Template Path (#49) (#51)

* Feature: Adding an option for custom template path

* restructure and small changes

- Removed unnecessary requirements
- Removed some unused imports
- Removed CUSTOM_RIG_TEMPLATES_PATH constant
- Fixed create new template issue of not adding the default metarig right away if selected starter for some reason isn't a Rigify default
- Moved get_rig_template_path() to utilities (second guessing this change)

---------

Co-authored-by: Alexander Dmitriev <dmit94alex@gmail.com>
Co-authored-by: JoshQuake <j.wake@hotmail.com>
---
 requirements.txt                             |  4 +--
 src/addons/ue2rigify/__init__.py             |  6 ++--
 src/addons/ue2rigify/constants.py            |  2 +-
 src/addons/ue2rigify/core/nodes.py           |  4 +--
 src/addons/ue2rigify/core/scene.py           | 16 ++++++++--
 src/addons/ue2rigify/core/templates.py       | 23 +++++++++------
 src/addons/ue2rigify/core/utilities.py       | 15 +++++++++-
 src/addons/ue2rigify/operators.py            | 16 ++++++++--
 src/addons/ue2rigify/properties.py           |  6 ++--
 src/addons/ue2rigify/ui/addon_preferences.py | 31 +++++++++++++++++---
 src/addons/ue2rigify/ui/view_3d.py           |  3 ++
 11 files changed, 97 insertions(+), 29 deletions(-)

diff --git a/requirements.txt b/requirements.txt
index cc0f56b9..0c15e219 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -7,11 +7,11 @@ numpy==1.26.4
 websocket==0.2.1
 python-dotenv==1.0.1
 
-# documenation 
+# documenation
 mkdocs<=1.5.3
 mkdocs-material<=9.5.13
 mkdocs-material-extensions<=1.3.1
 mkdocs-mermaid2-plugin<=1.1.1
 mkdocs-monorepo-plugin<=1.1.0
 termynal<=0.12.1
-mkdocstrings[python]<=1.9.0
\ No newline at end of file
+mkdocstrings[python]<=1.9.0
diff --git a/src/addons/ue2rigify/__init__.py b/src/addons/ue2rigify/__init__.py
index 04171fb4..53748105 100644
--- a/src/addons/ue2rigify/__init__.py
+++ b/src/addons/ue2rigify/__init__.py
@@ -45,8 +45,6 @@ def register():
     Registers the addon classes when the addon is enabled.
     """
 
-    templates.copy_default_templates()
-
     # reload the submodules
     if os.environ.get('UE2RIGIFY_DEV'):
         for module in modules:
@@ -58,6 +56,8 @@ def register():
     operators.register()
     nodes.register()
 
+    templates.copy_default_templates()
+
 
 def unregister():
     """
@@ -73,4 +73,4 @@ def unregister():
         properties.unregister()
     except Exception as e:
         print(f"Error un-registering UE2Rigify: \n{e}")
-        
+
diff --git a/src/addons/ue2rigify/constants.py b/src/addons/ue2rigify/constants.py
index 5c6f1919..4bf541c6 100644
--- a/src/addons/ue2rigify/constants.py
+++ b/src/addons/ue2rigify/constants.py
@@ -34,7 +34,7 @@ class Nodes:
 
 class Template:
     @staticmethod
-    def RIG_TEMPLATES_PATH():
+    def DEFAULT_RIG_TEMPLATES_PATH():
         if bpy.app.version[0] < 4:
             return os.path.join(tempfile.gettempdir(), ToolInfo.NAME.value, 'resources', 'rig_templates', 'b3_6')
         else:
diff --git a/src/addons/ue2rigify/core/nodes.py b/src/addons/ue2rigify/core/nodes.py
index 2316f483..a981a685 100644
--- a/src/addons/ue2rigify/core/nodes.py
+++ b/src/addons/ue2rigify/core/nodes.py
@@ -8,9 +8,9 @@
 from . import scene
 from . import utilities
 from ..ui import node_editor
-from ..constants import ToolInfo, Nodes, Modes, Rigify
+from ..constants import Nodes, Modes, Rigify
 from bpy.types import NodeTree, NodeSocket
-from nodeitems_utils import NodeCategory, NodeItem
+from nodeitems_utils import NodeItem
 
 addon_key_maps = []
 pie_menu_classes = []
diff --git a/src/addons/ue2rigify/core/scene.py b/src/addons/ue2rigify/core/scene.py
index e66ec1d4..561607fc 100644
--- a/src/addons/ue2rigify/core/scene.py
+++ b/src/addons/ue2rigify/core/scene.py
@@ -267,7 +267,7 @@ def set_meta_rig(self=None, context=None):
     else:
         remove_metarig()
         create_meta_rig(self, file_path=os.path.join(
-            Template.RIG_TEMPLATES_PATH(),
+            utilities.get_rig_template_path(),
             self.selected_starter_metarig_template,
             f'{Rigify.META_RIG_NAME}.py'
         ))
@@ -697,7 +697,17 @@ def create_starter_metarig_template(properties):
 
     # create the selected metarig
     operator = properties.selected_starter_metarig_template
-    exec(operator)
+    if operator.startswith('bpy.ops'):
+        exec(operator)
+    else:
+        # If selected starter isn't set nor a default Rigify template i.e. ue manny
+        if not operator:
+            operator = Template.DEFAULT_MALE_TEMPLATE
+        create_meta_rig(properties, file_path=os.path.join(
+            utilities.get_rig_template_path(),
+            operator,
+            f'{Rigify.META_RIG_NAME}.py'
+        ))
 
     # get the meta rig object
     meta_rig_object = bpy.data.objects.get(Rigify.META_RIG_NAME)
@@ -1449,7 +1459,7 @@ def edit_meta_rig_template(properties):
 
     :param object properties: The property group that contains variables that maintain the addon's correct state.
     """
-    # if not creating a new template, initiate the
+    # if not creating a new template, initiate the metarig
     if properties.selected_rig_template != 'create_new':
         meta_rig_object = create_meta_rig(properties)
     else:
diff --git a/src/addons/ue2rigify/core/templates.py b/src/addons/ue2rigify/core/templates.py
index 53fbb8e1..ce9c10ba 100644
--- a/src/addons/ue2rigify/core/templates.py
+++ b/src/addons/ue2rigify/core/templates.py
@@ -5,7 +5,6 @@
 import bpy
 import json
 import shutil
-import tempfile
 from mathutils import Color, Euler, Matrix, Quaternion, Vector
 
 from ..constants import Template, Modes, Rigify
@@ -26,7 +25,7 @@ def copy_default_templates():
     sub_folder = 'b3_6' if bpy.app.version[0] < 4 else 'b4_0'
     template_location = os.path.join(os.path.dirname(__file__), os.path.pardir, 'resources', 'rig_templates', sub_folder)
 
-    shutil.copytree(template_location, Template.RIG_TEMPLATES_PATH(), dirs_exist_ok=True)
+    shutil.copytree(template_location, utilities.get_rig_template_path(), dirs_exist_ok=True)
 
 
 def get_saved_node_data(properties):
@@ -189,7 +188,7 @@ def get_rig_templates(self=None, context=None, index_offset=0):
     :return list: A list of tuples that define the rig template enumeration.
     """
     rig_templates = []
-    rig_template_directories = next(os.walk(Template.RIG_TEMPLATES_PATH()))[1]
+    rig_template_directories = next(os.walk(utilities.get_rig_template_path()))[1]
 
     # ensure that the male and female template are first
     rig_templates.append((
@@ -235,7 +234,7 @@ def get_template_file_path(template_file_name, properties):
         template_name = re.sub(r'\W+', '_', properties.new_template_name).lower()
 
     return os.path.join(
-        Template.RIG_TEMPLATES_PATH(),
+        utilities.get_rig_template_path(),
         template_name,
         template_file_name
     )
@@ -286,7 +285,7 @@ def remove_template_folder(properties, template):
     properties.selected_mode = Modes.SOURCE.name
 
     # delete the selected rig template folder
-    selected_template_path = os.path.join(Template.RIG_TEMPLATES_PATH(), template)
+    selected_template_path = os.path.join(utilities.get_rig_template_path(), template)
 
     import logging
     logging.info(selected_template_path)
@@ -304,10 +303,10 @@ def create_template_folder(template_name, properties):
     :param object properties: The property group that contains variables that maintain the addon's correct state.
     """
     # remove non alpha numeric characters
-    template_name = re.sub(r'\W+', '_', template_name.strip()).lower()
+    template_name = parse_template_name(template_name)
 
     # create the template folder
-    template_path = os.path.join(Template.RIG_TEMPLATES_PATH(), template_name)
+    template_path = os.path.join(utilities.get_rig_template_path(), template_name)
     if not os.path.exists(template_path):
         os.makedirs(template_path, exist_ok=True)
 
@@ -396,7 +395,7 @@ def import_zip(zip_file_path, properties):
     """
     # get the template name and path from the zip file
     template_name = os.path.basename(zip_file_path).replace('.zip', '').lower()
-    template_folder_path = os.path.join((Template.RIG_TEMPLATES_PATH()), template_name)
+    template_folder_path = os.path.join(utilities.get_rig_template_path(), template_name)
 
     # create the template folder
     create_template_folder(template_name, properties)
@@ -404,6 +403,8 @@ def import_zip(zip_file_path, properties):
     # unpack the zip file into the new template folder
     shutil.unpack_archive(zip_file_path, template_folder_path, 'zip')
 
+    # select new template
+    properties.selected_rig_template = parse_template_name(template_name)
 
 # TODO move to a shared location and define as a generic zip export
 def export_zip(zip_file_path, properties):
@@ -417,7 +418,7 @@ def export_zip(zip_file_path, properties):
     no_extension_file_path = zip_file_path.replace('.zip', '')
 
     # zip up the folder and save it to the given path
-    template_folder_path = os.path.join(Template.RIG_TEMPLATES_PATH(), properties.selected_export_template)
+    template_folder_path = os.path.join(utilities.get_rig_template_path(), properties.selected_export_template)
     shutil.make_archive(no_extension_file_path, 'zip', template_folder_path)
 
 #
@@ -484,3 +485,7 @@ def safe_get_rig_templates(self, context):
     global _result_reference_get_rig_templates
     _result_reference_get_rig_templates = items
     return items
+
+def parse_template_name(template_name):
+    return re.sub(r'\W+', '_', template_name.strip()).lower()
+
diff --git a/src/addons/ue2rigify/core/utilities.py b/src/addons/ue2rigify/core/utilities.py
index 998fd756..e0b3c9c1 100644
--- a/src/addons/ue2rigify/core/utilities.py
+++ b/src/addons/ue2rigify/core/utilities.py
@@ -1,8 +1,9 @@
 # Copyright Epic Games, Inc. All Rights Reserved.
 import bpy
 import re
+import os
 from . import scene, templates
-from ..constants import Viewport, Modes, Rigify
+from ..constants import Viewport, Modes, Rigify, Template, ToolInfo
 from mathutils import Vector, Quaternion
 
 
@@ -1183,3 +1184,15 @@ def get_rigify_bone_operator(un_hashed_operator_name, bone_name, properties):
 
                 if output_bones and input_bones and ctrl_bones:
                     return f'bpy.ops.{operator}({output_bones}, {input_bones}, {ctrl_bones})'
+
+def get_rig_template_path():
+    preferences = bpy.context.preferences.addons.get(ToolInfo.NAME.value).preferences
+
+    template_path = preferences.custom_rig_template_path
+    # If custom_rig_template_path is empty, returns Temp folder
+    if template_path:
+        subpath = 'b3_6' if bpy.app.version[0] < 4 else 'b4_0'
+        template_path = os.path.join(template_path, subpath)
+    else:
+        template_path = Template.DEFAULT_RIG_TEMPLATES_PATH()
+    return template_path
\ No newline at end of file
diff --git a/src/addons/ue2rigify/operators.py b/src/addons/ue2rigify/operators.py
index 7e1efe18..8a29e61d 100644
--- a/src/addons/ue2rigify/operators.py
+++ b/src/addons/ue2rigify/operators.py
@@ -1,11 +1,22 @@
 # Copyright Epic Games, Inc. All Rights Reserved.
 
 import bpy
-from .constants import Modes, Rigify
+from .constants import Modes, Rigify, Template
 from .ui import exporter
 from .core import scene, nodes, templates, utilities
 from bpy_extras.io_utils import ImportHelper
 
+class SetDefault(bpy.types.Operator):
+    """Set a default control rig mode"""
+    bl_idname = "ue2rigify.set_default"
+    bl_label = "Default"
+
+    def execute(self, context):
+        properties = bpy.context.scene.ue2rigify
+        properties.selected_mode = Modes.SOURCE.name
+        properties.selected_rig_template = Template.DEFAULT_MALE_TEMPLATE
+        return {'FINISHED'}
+
 
 class ConvertToRigifyRig(bpy.types.Operator):
     """Convert the source rig to a control rig"""
@@ -294,7 +305,8 @@ def execute(self, context):
     ConstrainSourceToDeform,
     RemoveConstraints,
     SwitchModes,
-    NullOperator
+    NullOperator,
+    SetDefault
 ]
 
 
diff --git a/src/addons/ue2rigify/properties.py b/src/addons/ue2rigify/properties.py
index bcfb4412..63723446 100644
--- a/src/addons/ue2rigify/properties.py
+++ b/src/addons/ue2rigify/properties.py
@@ -57,7 +57,8 @@ class UE2RigifyProperties(bpy.types.PropertyGroup):
         name="Metarig",
         description=tool_tips.starter_metarig_template_tool_tip,
         items=templates.safe_get_starter_metarig_templates,
-        update=scene.set_meta_rig
+        update=scene.set_meta_rig,
+        default=None
     )
 
     selected_rig_template: bpy.props.EnumProperty(
@@ -65,7 +66,8 @@ class UE2RigifyProperties(bpy.types.PropertyGroup):
         description=tool_tips.rig_template_tool_tip,
         items=templates.safe_populate_templates_dropdown,
         options={'ANIMATABLE'},
-        update=templates.set_template
+        update=templates.set_template,
+        default=None
     )
 
     selected_mode: bpy.props.EnumProperty(
diff --git a/src/addons/ue2rigify/ui/addon_preferences.py b/src/addons/ue2rigify/ui/addon_preferences.py
index 9e3d5d25..828e1fd2 100644
--- a/src/addons/ue2rigify/ui/addon_preferences.py
+++ b/src/addons/ue2rigify/ui/addon_preferences.py
@@ -2,6 +2,7 @@
 
 import bpy
 from ..constants import ToolInfo
+from ..core import templates
 
 
 class Ue2RigifyAddonPreferences(bpy.types.AddonPreferences):
@@ -9,6 +10,28 @@ class Ue2RigifyAddonPreferences(bpy.types.AddonPreferences):
     This class subclasses the AddonPreferences class to create the addon preferences interface.
     """
     bl_idname = ToolInfo.NAME.value
+    
+    def get_custom_location(self):
+        # create key if doesn't exist then return
+        try:
+            self['custom_template_path']
+        except:
+            self['custom_template_path'] = ''
+        return self['custom_template_path']
+
+    def set_custom_location(self, value):
+        self['custom_template_path'] = value
+        # Create default templates at custom_rig_template_path
+        templates.copy_default_templates()
+
+    custom_rig_template_path: bpy.props.StringProperty(
+        name='Custom Templates folder',
+        description='The location where your rig templates will be stored, including default templates. Defaults to Temp folder if empty',
+        subtype='DIR_PATH',
+        default='',
+        get=get_custom_location,
+        set=set_custom_location
+    )
 
     def draw(self, context):
         """
@@ -19,6 +42,7 @@ def draw(self, context):
         """
         layout = self.layout
 
+        layout.prop(self, 'custom_rig_template_path')
         row = layout.row()
         row.operator('ue2rigify.import_rig_template', icon='IMPORT')
         row.operator('ue2rigify.export_rig_template', icon='EXPORT')
@@ -28,13 +52,12 @@ def register():
     """
     Registers the addon preferences when the addon is enabled.
     """
-    if not hasattr(bpy.types, Ue2RigifyAddonPreferences.bl_idname):
-        bpy.utils.register_class(Ue2RigifyAddonPreferences)
+    bpy.utils.register_class(Ue2RigifyAddonPreferences)
 
 
 def unregister():
     """
     Unregisters the addon preferences when the addon is disabled.
     """
-    if hasattr(bpy.types, Ue2RigifyAddonPreferences.bl_idname):
-        bpy.utils.unregister_class(Ue2RigifyAddonPreferences)
+    
+    bpy.utils.unregister_class(Ue2RigifyAddonPreferences)
diff --git a/src/addons/ue2rigify/ui/view_3d.py b/src/addons/ue2rigify/ui/view_3d.py
index c98a89a8..c6f70584 100644
--- a/src/addons/ue2rigify/ui/view_3d.py
+++ b/src/addons/ue2rigify/ui/view_3d.py
@@ -62,6 +62,9 @@ def draw(self, context):
             row = layout.row()
             row.prop(properties, 'selected_mode', text='')
 
+        if properties.selected_rig_template == '':
+            row.operator('ue2rigify.set_default')
+
         box = layout.box()
         row = box.row()
         row.label(text='Rig Template Editor', icon='TOOL_SETTINGS')