Rigify: implement a new reload process to support Reload Scripts properly.

Due to dynamically loaded rig modules, Rigify has to implement
reload by purging all non-core modules from memory, or class
inheritance checks won't work right after reload.

6 files changed, 135 insertions, 67 deletions

diff --git a/rigify/__init__.py b/rigify/__init__.py
index d487f265..803c19c0 100644
--- a/rigify/__init__.py
+++ b/rigify/__init__.py
@@ -29,24 +29,97 @@ bl_info = {
                 "Scripts/Rigging/Rigify",
     "category": "Rigging"}
 
+import importlib
+import sys
+import bpy
+import os
 
-if "bpy" in locals():
-    import importlib
-    # Don't reload base_rig or base_generate, because it would break issubclass checks,
-    # unless _all_ modules with classes inheriting from BaseRig are also reloaded.
-    importlib.reload(utils)
-    importlib.reload(rig_ui_template)
-    importlib.reload(feature_set_list)
-    importlib.reload(rig_lists)
-    importlib.reload(generate)
-    importlib.reload(ui)
-    importlib.reload(metarig_menu)
+
+# The order in which core modules of the addon are loaded and reloaded.
+# Modules not in this list are removed from memory upon reload.
+# With the sole exception of 'utils', modules must be listed in the
+# correct dependency order.
+initial_load_order = [
+    'utils',
+    'utils.errors',
+    'utils.misc',
+    'utils.rig',
+    'utils.naming',
+    'utils.bones',
+    'utils.collections',
+    'utils.layers',
+    'utils.widgets',
+    'utils.widgets_basic',
+    'utils.widgets_special',
+    'utils.mechanism',
+    'utils.animation',
+    'utils.metaclass',
+    'feature_sets',
+    'rigs',
+    'rigs.utils',
+    'base_rig',
+    'base_generate',
+    'feature_set_list',
+    'rig_lists',
+    'metarig_menu',
+    'rig_ui_template',
+    'generate',
+    'rot_mode',
+    'ui',
+]
+
+
+def get_loaded_modules():
+    prefix = __name__ + '.'
+    return [name for name in sys.modules if name.startswith(prefix)]
+
+def reload_modules():
+    fixed_modules = set(reload_list)
+
+    for name in get_loaded_modules():
+        if name not in fixed_modules:
+            del sys.modules[name]
+
+    for name in reload_list:
+        importlib.reload(sys.modules[name])
+
+def load_initial_modules():
+    load_list = [ __name__ + '.' + name for name in initial_load_order ]
+
+    for i, name in enumerate(load_list):
+        importlib.import_module(name)
+
+        module_list = get_loaded_modules()
+        expected_list = load_list[0 : max(11, i+1)]
+
+        if module_list != expected_list:
+            print('!!! RIGIFY: initial load order mismatch after '+name+' - expected: \n', expected_list, '\nGot:\n', module_list)
+
+    return load_list
+
+def load_rigs():
+    if not legacy_loaded:
+        rig_lists.get_internal_rigs()
+        metarig_menu.init_metarig_menu()
+
+
+if "reload_list" in locals():
+    reload_modules()
 else:
-    from . import (utils, base_rig, base_generate, rig_ui_template, feature_set_list, rig_lists, generate, ui, metarig_menu)
+    legacy_loaded = False
+
+    load_list = load_initial_modules()
+
+    from . import (base_rig, base_generate, rig_ui_template, feature_set_list, rig_lists, generate, ui, metarig_menu)
+
+    reload_list = reload_list_init = get_loaded_modules()
+
+    if reload_list != load_list:
+        print('!!! RIGIFY: initial load order mismatch - expected: \n', load_list, '\nGot:\n', reload_list)
+
+load_rigs()
+
 
-import bpy
-import sys
-import os
 from bpy.types import AddonPreferences
 from bpy.props import (
     BoolProperty,
@@ -65,16 +138,14 @@ class RigifyPreferences(AddonPreferences):
     bl_idname = __name__
 
     def update_legacy(self, context):
-        if self.legacy_mode:
+        global legacy_loaded, reload_list
 
-            if 'ui' in globals() and 'legacy' in str(globals()['ui']):    # already in legacy mode. needed when rigify is reloaded
+        if self.legacy_mode:
+            if legacy_loaded:    # already in legacy mode. needed when rigify is reloaded
                 return
             else:
-                rigify_dir = os.path.dirname(os.path.realpath(__file__))
-                if rigify_dir not in sys.path:
-                    sys.path.append(rigify_dir)
-
                 unregister()
+                reload_modules()
 
                 globals().pop('utils')
                 globals().pop('rig_lists')
@@ -82,14 +153,13 @@ class RigifyPreferences(AddonPreferences):
                 globals().pop('ui')
                 globals().pop('metarig_menu')
 
-                import legacy.utils
-                import legacy.rig_lists
-                import legacy.generate
-                import legacy.ui
-                import legacy.metarig_menu
+                from .legacy import utils, rig_lists, generate, ui, metarig_menu
 
                 print("ENTERING RIGIFY LEGACY\r\n")
 
+                legacy_loaded = True
+                reload_list += [ m.__name__ for m in [ legacy, utils, rig_lists, generate, ui, metarig_menu ] ]
+
                 globals()['utils'] = legacy.utils
                 globals()['rig_lists'] = legacy.rig_lists
                 globals()['generate'] = legacy.generate
@@ -99,13 +169,6 @@ class RigifyPreferences(AddonPreferences):
                 register()
 
         else:
-
-            rigify_dir = os.path.dirname(os.path.realpath(__file__))
-
-            if rigify_dir in sys.path:
-                id = sys.path.index(rigify_dir)
-                sys.path.pop(id)
-
             unregister()
 
             globals().pop('utils')
@@ -114,11 +177,7 @@ class RigifyPreferences(AddonPreferences):
             globals().pop('ui')
             globals().pop('metarig_menu')
 
-            from . import utils
-            from . import rig_lists
-            from . import generate
-            from . import ui
-            from . import metarig_menu
+            from . import utils, rig_lists, generate, ui, metarig_menu
 
             print("EXIT RIGIFY LEGACY\r\n")
 
@@ -128,6 +187,12 @@ class RigifyPreferences(AddonPreferences):
             globals()['ui'] = ui
             globals()['metarig_menu'] = metarig_menu
 
+            legacy_loaded = False
+            reload_list = reload_list_init
+            reload_modules()
+
+            load_rigs()
+
             register()
 
     def update_external_rigs(self, force=False):
@@ -492,7 +557,7 @@ def register():
         description="Transfer selected bones only", default=True)
 
     # Update legacy on restart or reload.
-    if (ui and 'legacy' in str(ui)) or bpy.context.preferences.addons['rigify'].preferences.legacy_mode:
+    if legacy_loaded or bpy.context.preferences.addons['rigify'].preferences.legacy_mode:
         bpy.context.preferences.addons['rigify'].preferences.legacy_mode = True
 
     bpy.context.preferences.addons['rigify'].preferences.update_external_rigs()
diff --git a/rigify/metarig_menu.py b/rigify/metarig_menu.py
index dc583a67..48da1e89 100644
--- a/rigify/metarig_menu.py
+++ b/rigify/metarig_menu.py
@@ -26,7 +26,8 @@ from collections import defaultdict
 
 import bpy
 
-from . import utils
+from .utils.rig import METARIG_DIR, get_resource
+
 from . import feature_set_list
 
 
@@ -70,11 +71,11 @@ def get_metarigs(metarigs, base_dir, base_path, *, path=[], nested=False):
         elif f.endswith(".py"):
             # Check straight-up python files
             f = f[:-3]
-            module = utils.get_resource('.'.join([*base_path, *path, f]))
+            module = get_resource('.'.join([*base_path, *path, f]))
             if nested:
                 metarigs[f] = module
             else:
-                metarigs[utils.METARIG_DIR][f] = module
+                metarigs[METARIG_DIR][f] = module
 
 
 def make_metarig_add_execute(m):
@@ -120,7 +121,7 @@ def get_internal_metarigs():
     BASE_RIGIFY_DIR = os.path.dirname(__file__)
     BASE_RIGIFY_PATH = __name__.split('.')[:-1]
 
-    get_metarigs(metarigs, os.path.join(BASE_RIGIFY_DIR, utils.METARIG_DIR), [*BASE_RIGIFY_PATH, utils.METARIG_DIR])
+    get_metarigs(metarigs, os.path.join(BASE_RIGIFY_DIR, METARIG_DIR), [*BASE_RIGIFY_PATH, METARIG_DIR])
 
 def infinite_defaultdict():
     return defaultdict(infinite_defaultdict)
@@ -130,8 +131,6 @@ metarig_ops = {}
 armature_submenus = []
 menu_funcs = []
 
-get_internal_metarigs()
-
 def create_metarig_ops(dic=metarigs):
     """Create metarig add Operators"""
     for metarig_category in dic:
@@ -154,7 +153,7 @@ def create_metarig_ops(dic=metarigs):
 
 def create_menu_funcs():
     global menu_funcs
-    for mop, name in metarig_ops[utils.METARIG_DIR]:
+    for mop, name in metarig_ops[METARIG_DIR]:
         text = capwords(name.replace("_", " ")) + " (Meta-Rig)"
         menu_funcs += [make_metarig_menu_func(mop.bl_idname, text)]
 
@@ -167,7 +166,7 @@ def create_armature_submenus(dic=metarigs):
         if metarig_category == "external":
             create_armature_submenus(dic=metarigs["external"])
             continue
-        if metarig_category == utils.METARIG_DIR:
+        if metarig_category == METARIG_DIR:
             continue
 
         armature_submenus.append(type('Class_' + metarig_category + '_submenu', (ArmatureSubMenu,), {}))
@@ -180,10 +179,11 @@ def create_armature_submenus(dic=metarigs):
             arm_sub = next((e for e in armature_submenus if e.bl_label == metarig_category + ' (submenu)'), '')
             arm_sub.operators.append((mop.bl_idname, name,))
 
-create_metarig_ops()
-create_menu_funcs()
-create_armature_submenus()
-
+def init_metarig_menu():
+    get_internal_metarigs()
+    create_metarig_ops()
+    create_menu_funcs()
+    create_armature_submenus()
 
 ### Registering ###
 
@@ -224,7 +224,7 @@ def get_external_metarigs(set_list):
 
     for feature_set in set_list:
         try:
-            base_dir, base_path = feature_set_list.get_dir_path(feature_set, utils.METARIG_DIR)
+            base_dir, base_path = feature_set_list.get_dir_path(feature_set, METARIG_DIR)
 
             get_metarigs(metarigs['external'], base_dir, base_path)
         except Exception:
diff --git a/rigify/rig_lists.py b/rigify/rig_lists.py
index 018bbbac..49cc9545 100644
--- a/rigify/rig_lists.py
+++ b/rigify/rig_lists.py
@@ -20,7 +20,8 @@ import os
 import traceback
 import importlib
 
-from . import utils
+from .utils.rig import RIG_DIR
+
 from . import feature_set_list
 
 
@@ -76,15 +77,16 @@ def get_rigs(base_dir, base_path, *, path=[], feature_set=feature_set_list.DEFAU
 
 
 # Public variables
+rigs = {}
+implementation_rigs = {}
+
 def get_internal_rigs():
+    global rigs, implementation_rigs
+
     BASE_RIGIFY_DIR = os.path.dirname(__file__)
     BASE_RIGIFY_PATH = __name__.split('.')[:-1]
 
-    return get_rigs(os.path.join(BASE_RIGIFY_DIR, utils.RIG_DIR), [*BASE_RIGIFY_PATH, utils.RIG_DIR])
-
-
-rigs, implementation_rigs = get_internal_rigs()
-
+    rigs, implementation_rigs = get_rigs(os.path.join(BASE_RIGIFY_DIR, RIG_DIR), [*BASE_RIGIFY_PATH, RIG_DIR])
 
 def get_external_rigs(set_list):
     # Clear and fill rigify rigs and implementation rigs public variables
@@ -97,7 +99,7 @@ def get_external_rigs(set_list):
     # Get external rigs
     for feature_set in set_list:
         try:
-            base_dir, base_path = feature_set_list.get_dir_path(feature_set, utils.RIG_DIR)
+            base_dir, base_path = feature_set_list.get_dir_path(feature_set, RIG_DIR)
 
             external_rigs, external_impl_rigs = get_rigs(base_dir, base_path, feature_set=feature_set)
         except Exception:
diff --git a/rigify/rigs/utils.py b/rigify/rigs/utils.py
index 12dcbb26..89873ca4 100644
--- a/rigify/rigs/utils.py
+++ b/rigify/rigs/utils.py
@@ -1,4 +1,4 @@
-from ..utils import connected_children_names
+from ..utils.rig import connected_children_names
 from ..utils.naming import strip_mch, strip_org, make_mechanism_name
 import re
 
diff --git a/rigify/ui.py b/rigify/ui.py
index 2f858925..dacbd0a2 100644
--- a/rigify/ui.py
+++ b/rigify/ui.py
@@ -28,10 +28,12 @@ from bpy.props import (
 
 from mathutils import Color
 
-from .utils import MetarigError
-from .utils import write_metarig, write_widget
-from .utils import unique_name
-from .utils import upgradeMetarigTypes, outdated_types
+from .utils.errors import MetarigError
+from .utils.rig import write_metarig
+from .utils.widgets import write_widget
+from .utils.naming import unique_name
+from .utils.rig import upgradeMetarigTypes, outdated_types
+
 from .rigs.utils import get_limb_generated_names
 
 from .utils.animation import get_keyed_frames_in_range, bones_in_frame, overwrite_prop_animation
@@ -767,9 +769,6 @@ class Generate(bpy.types.Operator):
     bl_description = 'Generates a rig from the active metarig armature'
 
     def execute(self, context):
-        import importlib
-        importlib.reload(generate)
-
         try:
             generate.generate_rig(context, context.object)
         except MetarigError as rig_exception:
diff --git a/rigify/utils/__init__.py b/rigify/utils/__init__.py
index f45acded..f2444d1c 100644
--- a/rigify/utils/__init__.py
+++ b/rigify/utils/__init__.py
@@ -2,6 +2,8 @@
 # that expects a single utils.py file. New code should import directly from
 # the modules that contain the utilities. Also, don't add more imports here.
 
+from . import errors, misc, rig, naming, bones, collections, layers, widgets, widgets_basic, widgets_special
+
 from .errors import MetarigError
 
 from .misc import angle_on_plane, linsrgb_to_srgb, gamma_correct, copy_attributes