Rigify: Fix collections/layers

Now all the widget elements are in the Widgets collection.
And the collection is hidden for the viewport.

There is one thing still not fully working, but it is a bug in Blender I
believe. Basically not all the custom bones are shown until you go to an
object that has its mesh and toggle in and out of edit mode.

3 files changed, 47 insertions, 7 deletions

diff --git a/rigify/legacy/generate.py b/rigify/legacy/generate.py
index 4cc725b1..9a39e333 100644
--- a/rigify/legacy/generate.py
+++ b/rigify/legacy/generate.py
@@ -28,7 +28,7 @@ from rna_prop_ui import rna_idprop_ui_prop_get
 from .utils import MetarigError, new_bone, get_rig_type
 from .utils import ORG_PREFIX, MCH_PREFIX, DEF_PREFIX, WGT_PREFIX, ROOT_NAME, make_original_name
 from .utils import RIG_DIR
-from .utils import create_root_widget
+from .utils import create_root_widget, ensure_widget_collection
 from .utils import random_id
 from .utils import copy_attributes
 from .rig_ui_template import UI_SLIDERS, layers_ui, UI_REGISTER
@@ -73,6 +73,7 @@ def generate_rig(context, metarig):
     scene = context.scene
     view_layer = context.view_layer
     collection = scene.collection
+    layer_collection = context.layer_collection
 
     #------------------------------------------
     # Create/find the rig object and set it up
@@ -355,7 +356,10 @@ def generate_rig(context, metarig):
         if obj.data.bones[bone].name.startswith(DEF_PREFIX):
             obj.data.bones[bone].layers = DEF_LAYER
 
-    # Create root bone widget
+    # Create/find widge collection
+    ensure_widget_collection(context)
+
+    #  Create root bone widget
     create_root_widget(obj, "root")
 
     # Assign shapes to bones
@@ -429,6 +433,10 @@ def generate_rig(context, metarig):
     metarig.data.pose_position = rest_backup
     obj.data.pose_position = 'POSE'
 
+    #----------------------------------
+    # Restore active collection
+    view_layer.collections.active = layer_collection
+
 
 def get_bone_rigs(obj, bone_name, halt_on_missing=False):
     """ Fetch all the rigs specified on a bone.
diff --git a/rigify/legacy/rigs/pitchipoy/super_widgets.py b/rigify/legacy/rigs/pitchipoy/super_widgets.py
index 72384a7c..f442c590 100644
--- a/rigify/legacy/rigs/pitchipoy/super_widgets.py
+++ b/rigify/legacy/rigs/pitchipoy/super_widgets.py
@@ -3,7 +3,6 @@ import importlib
 import importlib
 from   ...utils import create_widget
 
-WGT_LAYERS  = [x == 19 for x in range(0, 20)]  # Widgets go on the last scene layer.
 MODULE_NAME = "super_widgets"  # Windows/Mac blender is weird, so __package__ doesn't work
 
 
diff --git a/rigify/legacy/utils.py b/rigify/legacy/utils.py
index f70a9967..0f3a092e 100644
--- a/rigify/legacy/utils.py
+++ b/rigify/legacy/utils.py
@@ -36,8 +36,6 @@ DEF_PREFIX = "DEF-"  # Prefix of deformation bones.
 WGT_PREFIX = "WGT-"  # Prefix for widget objects
 ROOT_NAME = "root"   # Name of the root bone.
 
-WGT_LAYERS = [x == 19 for x in range(0, 20)]  # Widgets go on the last scene layer.
-
 MODULE_NAME = "rigify"  # Windows/Mac blender is weird, so __package__ doesn't work
 
 
@@ -397,6 +395,7 @@ def create_widget(rig, bone_name, bone_transform_name=None):
 
     obj_name = WGT_PREFIX + bone_name
     scene = bpy.context.scene
+    collection = bpy.context.collection
 
     # Check if it already exists in the scene
     if obj_name in scene.objects:
@@ -420,8 +419,6 @@ def create_widget(rig, bone_name, bone_transform_name=None):
 
         # Move object to bone position and set layers
         obj_to_bone(obj, rig, bone_transform_name)
-        # TODO move colleciton to all the WGT_LAYERS collections
-        # obj.layers = WGT_LAYERS
 
         return obj
 
@@ -938,3 +935,39 @@ def random_id(length=8):
         text += random.choice(chars)
     text += str(hex(int(time.time())))[2:][-tlength:].rjust(tlength, '0')[::-1]
     return text
+
+
+def get_layer_collection_from_collection(children, collection):
+    for layer_collection in children:
+        if collection == layer_collection.collection:
+            return layer_collection
+
+        # go recursive
+        layer_collection = get_layer_collection_from_collection(layer_collection.children, collection)
+        if layer_collection:
+            return layer_collection
+
+
+def ensure_widget_collection(context):
+    wgts_collection_name = "Widgets"
+
+    view_layer = context.view_layer
+    layer_collection = bpy.context.layer_collection
+    collection = layer_collection.collection
+
+    widget_collection = bpy.data.collections.get(wgts_collection_name)
+    if not widget_collection:
+        # ------------------------------------------
+        # Create the widget collection
+        widget_collection = bpy.data.collections.new(wgts_collection_name)
+        widget_collection.hide_viewport = True
+        widget_collection.hide_render = True
+
+        collection.children.link(widget_collection)
+        widget_layer_collection = [c for c in layer_collection.children if c.collection == widget_collection][0]
+    else:
+        widget_layer_collection = get_layer_collection_from_collection(view_layer.collections, widget_collection)
+
+    # Make the widget the active collection for the upcoming added (widget) objects
+    view_layer.collections.active = widget_layer_collection
+    return widget_collection