Rigify: add utilities for working with rig instance parent chains.

2 files changed, 37 insertions, 19 deletions

diff --git a/rigify/utils/rig.py b/rigify/utils/rig.py
index fa55c1aa..7969a34f 100644
--- a/rigify/utils/rig.py
+++ b/rigify/utils/rig.py
@@ -91,6 +91,36 @@ def upgradeMetarigTypes(metarig, revert=False):
 # Misc
 #=============================================
 
+
+def rig_is_child(rig, parent, *, strict=False):
+    """
+    Checks if the rig is a child of the parent.
+    Unless strict is True, returns true if the rig and parent are the same.
+    """
+    if parent is None:
+        return True
+
+    if rig and strict:
+        rig = rig.rigify_parent
+
+    while rig:
+        if rig is parent:
+            return True
+
+        rig = rig.rigify_parent
+
+    return False
+
+
+def get_parent_rigs(rig):
+    """Returns a list containing the rig and all of its parents."""
+    result = []
+    while rig:
+        result.append(rig)
+        rig = rig.rigify_parent
+    return result
+
+
 def get_resource(resource_name):
     """ Fetches a rig module by name, and returns it.
     """
diff --git a/rigify/utils/switch_parent.py b/rigify/utils/switch_parent.py
index 1a0e81fe..405008b9 100644
--- a/rigify/utils/switch_parent.py
+++ b/rigify/utils/switch_parent.py
@@ -8,6 +8,7 @@ import json
 from .errors import MetarigError
 from .naming import strip_prefix, make_derived_name
 from .mechanism import MechanismUtilityMixin
+from .rig import rig_is_child
 from .misc import map_list, map_apply, force_lazy
 
 from ..base_rig import *
@@ -17,19 +18,6 @@ from collections import defaultdict
 from itertools import count, repeat, chain
 
 
-def _rig_is_child(rig, parent):
-    if parent is None:
-        return True
-
-    while rig:
-        if rig is parent:
-            return True
-
-        rig = rig.rigify_parent
-
-    return False
-
-
 class SwitchParentBuilder(GeneratorPlugin, MechanismUtilityMixin):
     """
     Implements centralized generation of switchable parent mechanisms.
@@ -71,8 +59,8 @@ class SwitchParentBuilder(GeneratorPlugin, MechanismUtilityMixin):
 
         assert not self.frozen
         assert isinstance(bone, str) or callable(bone)
-        assert callable(bone) or _rig_is_child(rig, self.generator.bone_owners[bone])
-        assert _rig_is_child(rig, inject_into)
+        assert callable(bone) or rig_is_child(rig, self.generator.bone_owners[bone])
+        assert rig_is_child(rig, inject_into)
 
         real_rig = rig
 
@@ -182,7 +170,7 @@ class SwitchParentBuilder(GeneratorPlugin, MechanismUtilityMixin):
 
     def assign_child_options(self, child, options):
         if 'context_rig' in options:
-            assert _rig_is_child(child['rig'], options['context_rig'])
+            assert rig_is_child(child['rig'], options['context_rig'])
 
         for name, value in options.items():
             if name not in self.child_option_table:
@@ -216,7 +204,7 @@ class SwitchParentBuilder(GeneratorPlugin, MechanismUtilityMixin):
 
                 # Exclude injected parents
                 if parent['real_rig'] is not parent_rig:
-                    if _rig_is_child(parent_rig, child_rig):
+                    if rig_is_child(parent_rig, child_rig):
                         continue
 
                 if parent['rig'] is child_rig:
@@ -224,11 +212,11 @@ class SwitchParentBuilder(GeneratorPlugin, MechanismUtilityMixin):
                         continue
                 elif parent['is_global'] and not child['ignore_global']:
                     # Can't use parents from own children, even if global (cycle risk)
-                    if _rig_is_child(parent_rig, child_rig):
+                    if rig_is_child(parent_rig, child_rig):
                         continue
                 else:
                     # Required to be a child of the parent's rig
-                    if not _rig_is_child(child_rig, parent_rig):
+                    if not rig_is_child(child_rig, parent_rig):
                         continue
 
                 parent['used'] = True