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.

6 files changed, 96 insertions, 36 deletions

diff --git a/rigify/generate.py b/rigify/generate.py
index 9d0d6051..948ece11 100644
--- a/rigify/generate.py
+++ b/rigify/generate.py
@@ -29,6 +29,7 @@ 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 ensure_widget_collection
 from .utils import random_id
 from .utils import copy_attributes
 from .utils import gamma_correct
@@ -40,7 +41,6 @@ ORG_LAYER = [n == 31 for n in range(0, 32)]  # Armature layer that original bone
 MCH_LAYER = [n == 30 for n in range(0, 32)]  # Armature layer that mechanism bones should be moved to.
 DEF_LAYER = [n == 29 for n in range(0, 32)]  # Armature layer that deformation bones should be moved to.
 ROOT_LAYER = [n == 28 for n in range(0, 32)]  # Armature layer that root bone should be moved to.
-WGT_LAYERS = [x == 19 for x in range(0, 20)]  # Widgets go on the last scene layer.
 
 
 class Timer:
@@ -73,8 +73,10 @@ def generate_rig(context, metarig):
 
     scene = context.scene
     view_layer = context.view_layer
-    collection = scene.collection
+    collection = context.collection
+    layer_collection = context.layer_collection
     id_store = context.window_manager
+
     #------------------------------------------
     # Create/find the rig object and set it up
 
@@ -122,21 +124,9 @@ def generate_rig(context, metarig):
     wgts_group_name = "WGTS_" + (rig_old_name or obj.name)
     if wgts_group_name in scene.objects and id_store.rigify_force_widget_update:
         bpy.ops.object.select_all(action='DESELECT')
-        # TODO handle collections visibility WGT_LAYERS
-        """
-        for i, lyr in enumerate(WGT_LAYERS):
-            if lyr:
-                context.scene.layers[i] = True
-        """
         for wgt in bpy.data.objects[wgts_group_name].children:
             wgt.select_set('SELECT')
         bpy.ops.object.delete(use_global=False)
-        # TODO handle collections visibility WGT_LAYERS
-        """
-        for i, lyr in enumerate(WGT_LAYERS):
-            if lyr:
-                context.scene.layers[i] = False
-        """
         if rig_old_name:
             bpy.data.objects[wgts_group_name].name = "WGTS_" + obj.name
 
@@ -311,7 +301,8 @@ def generate_rig(context, metarig):
     rna_idprop_ui_prop_get(obj.data, "rig_id", create=True)
     obj.data["rig_id"] = rig_id
 
-    t.tick("Create root bone: ")
+    # Create/find widge collection
+    widget_collection = ensure_widget_collection(context)
 
     # Create Group widget
     # wgts_group_name = "WGTS"
@@ -321,22 +312,14 @@ def generate_rig(context, metarig):
             bpy.data.objects.remove(bpy.data.objects[wgts_group_name])
         mesh = bpy.data.meshes.new(wgts_group_name)
         wgts_obj = bpy.data.objects.new(wgts_group_name, mesh)
-        collection.objects.link(wgts_obj)
-        # TODO MOVE objects to all WGT_LAYERS collection
-        #wgts_obj.layers = WGT_LAYERS
+        widget_collection.objects.link(wgts_obj)
         t.tick("Create main WGTS: ")
     #
     # if id_store.rigify_generate_mode == 'new':
     #     bpy.ops.object.select_all(action='DESELECT')
     #     for wgt in bpy.data.objects[wgts_group_name].children:
     #         wgt.select_set('SELECT')
-    #     for i, lyr in enumerate(WGT_LAYERS):
-    #         if lyr:
-    #             context.scene.layers[i] = True
     #     bpy.ops.object.make_single_user(obdata=True)
-    #     for i, lyr in enumerate(WGT_LAYERS):
-    #         if lyr:
-    #             context.scene.layers[i] = False
 
     #----------------------------------
     try:
@@ -536,6 +519,11 @@ def generate_rig(context, metarig):
             child.parent_bone = sub_parent
             child.matrix_world = mat
 
+    #----------------------------------
+    # Restore active collection
+    view_layer.collections.active = layer_collection
+
+
 def create_selection_sets(obj, metarig):
 
     # Check if selection sets addon is installed
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
diff --git a/rigify/rigs/widgets.py b/rigify/rigs/widgets.py
index 8461d82a..aebe7139 100644
--- a/rigify/rigs/widgets.py
+++ b/rigify/rigs/widgets.py
@@ -4,7 +4,6 @@ import importlib
 from mathutils import Matrix
 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/utils.py b/rigify/utils.py
index 2d5ac3f5..23596b3c 100644
--- a/rigify/utils.py
+++ b/rigify/utils.py
@@ -38,8 +38,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
 
 outdated_types = {"pitchipoy.limbs.super_limb": "limbs.super_limb",
@@ -440,7 +438,7 @@ def create_widget(rig, bone_name, bone_transform_name=None):
 
     obj_name = WGT_PREFIX + rig.name + '_' + bone_name
     scene = bpy.context.scene
-    collection = scene.collection #TODO WGT_LAYERS
+    collection = bpy.context.collection
     id_store = bpy.context.window_manager
 
     # Check if it already exists in the scene
@@ -468,7 +466,6 @@ def create_widget(rig, bone_name, bone_transform_name=None):
         wgts_group_name = 'WGTS_' + rig.name
         if wgts_group_name in bpy.data.objects.keys():
             obj.parent = bpy.data.objects[wgts_group_name]
-        #obj.layers = WGT_LAYERS #TODO WGT_LAYERS
 
         return obj
 
@@ -1263,3 +1260,39 @@ def overwrite_prop_animation(rig, bone, prop_name, value, frames):
     for kp in curve.keyframe_points:
         if kp.co[0] in 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
+
+        # 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