diff options
-rw-r--r-- | rigify/generate.py | 16 | ||||
-rw-r--r-- | rigify/utils/collections.py | 26 |
2 files changed, 40 insertions, 2 deletions
diff --git a/rigify/generate.py b/rigify/generate.py index 7eff06f3..2f2e3655 100644 --- a/rigify/generate.py +++ b/rigify/generate.py @@ -29,7 +29,7 @@ from collections import OrderedDict from .utils import MetarigError, new_bone from .utils import MCH_PREFIX, DEF_PREFIX, WGT_PREFIX, ROOT_NAME, make_original_name from .utils import create_root_widget -from .utils.collections import ensure_widget_collection +from .utils.collections import ensure_widget_collection, list_layer_collections, filter_layer_collections_by_object from .utils import random_id from .utils import copy_attributes from .utils import gamma_correct @@ -73,10 +73,17 @@ def generate_rig(context, metarig): scene = context.scene view_layer = context.view_layer - collection = context.collection layer_collection = context.layer_collection id_store = context.window_manager + usable_collections = list_layer_collections(view_layer.layer_collection, selectable=True) + + if layer_collection not in usable_collections: + metarig_collections = filter_layer_collections_by_object(usable_collections, metarig) + layer_collection = (metarig_collections + [view_layer.layer_collection])[0] + + collection = layer_collection.collection + #------------------------------------------ # Create/find the rig object and set it up @@ -96,6 +103,11 @@ def generate_rig(context, metarig): obj = scene.objects[name] rig_old_name = name obj.name = rig_new_name or name + + rig_collections = filter_layer_collections_by_object(usable_collections, obj) + layer_collection = (rig_collections + [layer_collection])[0] + collection = layer_collection.collection + except KeyError: rig_old_name = name name = rig_new_name or name diff --git a/rigify/utils/collections.py b/rigify/utils/collections.py index 25596905..5682ec64 100644 --- a/rigify/utils/collections.py +++ b/rigify/utils/collections.py @@ -39,6 +39,32 @@ def find_layer_collection_by_collection(layer_collection, collection): return layer_collection +def list_layer_collections(layer_collection, visible=False, selectable=False): + """Returns a list of the collection and its children, with optional filtering by settings.""" + + if layer_collection.exclude: + return [] + + collection = layer_collection.collection + is_visible = not (layer_collection.hide_viewport or collection.hide_viewport) + is_selectable = is_visible and not collection.hide_select + + if (selectable and not is_selectable) or (visible and not is_visible): + return [] + + found = [layer_collection] + + for child in layer_collection.children: + found += list_layer_collections(child, visible, selectable) + + return found + + +def filter_layer_collections_by_object(layer_collections, obj): + """Returns a subset of collections that contain the given object.""" + return [lc for lc in layer_collections if obj in lc.collection.objects.values()] + + def ensure_widget_collection(context): wgts_collection_name = "Widgets" |