Rigify: fix collection handling.

- Legacy mode didn't add objects to proper collections.
- Things broke if the Widgets collection already existed
  in another scene - linking the collection should be a
  separate step from creating it.

3 files changed, 28 insertions, 24 deletions

diff --git a/rigify/legacy/generate.py b/rigify/legacy/generate.py
index 4c571c3f..476cdbe8 100644
--- a/rigify/legacy/generate.py
+++ b/rigify/legacy/generate.py
@@ -72,7 +72,7 @@ def generate_rig(context, metarig):
 
     scene = context.scene
     view_layer = context.view_layer
-    collection = scene.collection
+    collection = context.collection
     layer_collection = context.layer_collection
 
     #------------------------------------------
@@ -265,6 +265,9 @@ def generate_rig(context, metarig):
     rna_idprop_ui_prop_get(obj.data, "rig_id", create=True)
     obj.data["rig_id"] = rig_id
 
+    # Create/find widget collection
+    ensure_widget_collection(context)
+
     t.tick("Create root bone: ")
     #----------------------------------
     try:
@@ -356,9 +359,6 @@ def generate_rig(context, metarig):
         if obj.data.bones[bone].name.startswith(DEF_PREFIX):
             obj.data.bones[bone].layers = DEF_LAYER
 
-    # Create/find widge collection
-    ensure_widget_collection(context)
-
     #  Create root bone widget
     create_root_widget(obj, "root")
 
diff --git a/rigify/legacy/utils.py b/rigify/legacy/utils.py
index 632dcb26..f74a7abf 100644
--- a/rigify/legacy/utils.py
+++ b/rigify/legacy/utils.py
@@ -937,13 +937,13 @@ def random_id(length=8):
     return text
 
 
-def get_layer_collection_from_collection(children, collection):
-    for layer_collection in children:
-        if collection == layer_collection.collection:
-            return layer_collection
+def find_layer_collection_by_collection(layer_collection, collection):
+    if collection == layer_collection.collection:
+        return layer_collection
 
-        # go recursive
-        layer_collection = get_layer_collection_from_collection(layer_collection.children, collection)
+    # go recursive
+    for child in layer_collection.children:
+        layer_collection = find_layer_collection_by_collection(child, collection)
         if layer_collection:
             return layer_collection
 
@@ -963,12 +963,14 @@ def ensure_widget_collection(context):
         widget_collection.hide_viewport = True
         widget_collection.hide_render = True
 
+        widget_layer_collection = None
+    else:
+        widget_layer_collection = find_layer_collection_by_collection(view_layer.layer_collection, widget_collection)
+
+    if not widget_layer_collection:
+        # Add the widget collection to the tree
         collection.children.link(widget_collection)
         widget_layer_collection = [c for c in layer_collection.children if c.collection == widget_collection][0]
-    elif widget_collection == view_layer.layer_collection.collection:
-        widget_layer_collection = view_layer.layer_collection
-    else:
-        widget_layer_collection = get_layer_collection_from_collection(view_layer.layer_collection.children, widget_collection)
 
     # Make the widget the active collection for the upcoming added (widget) objects
     view_layer.active_layer_collection = widget_layer_collection
diff --git a/rigify/utils.py b/rigify/utils.py
index 20dae315..73b64112 100644
--- a/rigify/utils.py
+++ b/rigify/utils.py
@@ -1262,13 +1262,13 @@ def overwrite_prop_animation(rig, bone, prop_name, value, frames):
             kp.co[1] = value
 
 
-def get_layer_collection_from_collection(children, collection):
-    for layer_collection in children:
-        if collection == layer_collection.collection:
-            return layer_collection
+def find_layer_collection_by_collection(layer_collection, collection):
+    if collection == layer_collection.collection:
+        return layer_collection
 
-        # go recursive
-        layer_collection = get_layer_collection_from_collection(layer_collection.children, collection)
+    # go recursive
+    for child in layer_collection.children:
+        layer_collection = find_layer_collection_by_collection(child, collection)
         if layer_collection:
             return layer_collection
 
@@ -1288,12 +1288,14 @@ def ensure_widget_collection(context):
         widget_collection.hide_viewport = True
         widget_collection.hide_render = True
 
+        widget_layer_collection = None
+    else:
+        widget_layer_collection = find_layer_collection_by_collection(view_layer.layer_collection, widget_collection)
+
+    if not widget_layer_collection:
+        # Add the widget collection to the tree
         collection.children.link(widget_collection)
         widget_layer_collection = [c for c in layer_collection.children if c.collection == widget_collection][0]
-    elif widget_collection == view_layer.layer_collection.collection:
-        widget_layer_collection = view_layer.layer_collection
-    else:
-        widget_layer_collection = get_layer_collection_from_collection(view_layer.layer_collection.children, widget_collection)
 
     # Make the widget the active collection for the upcoming added (widget) objects
     view_layer.active_layer_collection = widget_layer_collection