rigify/generate.py

   1 # SPDX-FileCopyrightText: 2010-2022 Blender Foundation
   2 #
   3 # SPDX-License-Identifier: GPL-2.0-or-later
   4
   5 import bpy
   6 import re
   7 import time
   8
   9 from typing import Optional, TYPE_CHECKING
  10
  11 from .utils.errors import MetarigError
  12 from .utils.bones import new_bone
  13 from .utils.layers import (ORG_COLLECTION, MCH_COLLECTION, DEF_COLLECTION, ROOT_COLLECTION, set_bone_layers,
  14                            validate_collection_references)
  15 from .utils.naming import (ORG_PREFIX, MCH_PREFIX, DEF_PREFIX, ROOT_NAME, make_original_name,
  16                            change_name_side, get_name_side, Side)
  17 from .utils.widgets import WGT_PREFIX, WGT_GROUP_PREFIX
  18 from .utils.widgets_special import create_root_widget
  19 from .utils.mechanism import refresh_all_drivers
  20 from .utils.misc import select_object, ArmatureObject, verify_armature_obj, choose_next_uid, flatten_children,\
  21     flatten_parents
  22 from .utils.collections import (ensure_collection, list_layer_collections,
  23                                 filter_layer_collections_by_object)
  24 from .utils.rig import get_rigify_type, get_rigify_target_rig,\
  25     get_rigify_rig_basename, get_rigify_force_widget_update, get_rigify_finalize_script,\
  26     get_rigify_mirror_widgets, get_rigify_colors
  27 from .utils.action_layers import ActionLayerBuilder
  28 from .utils.objects import ArtifactManager
  29
  30 from . import base_generate
  31 from . import rig_ui_template
  32 from . import rig_lists
  33
  34 if TYPE_CHECKING:
  35     from . import RigifyColorSet
  36
  37
  38 RIG_MODULE = "rigs"
  39
  40
  41 class Timer:
  42     def __init__(self):
  43         self.time_val = time.time()
  44
  45     def tick(self, string):
  46         t = time.time()
  47         print(string + "%.3f" % (t - self.time_val))
  48         self.time_val = t
  49
  50
  51 class Generator(base_generate.BaseGenerator):
  52     usable_collections: list[bpy.types.LayerCollection]
  53     action_layers: ActionLayerBuilder
  54
  55     def __init__(self, context, metarig):
  56         super().__init__(context, metarig)
  57
  58         self.id_store = context.window_manager
  59         self.saved_visible_layers = {}
  60
  61     def find_rig_class(self, rig_type):
  62         rig_module = rig_lists.rigs[rig_type]["module"]
  63
  64         return rig_module.Rig
  65
  66     def __switch_to_usable_collection(self, obj, fallback=False):
  67         collections = filter_layer_collections_by_object(self.usable_collections, obj)
  68
  69         if collections:
  70             self.layer_collection = collections[0]
  71         elif fallback:
  72             self.layer_collection = self.view_layer.layer_collection
  73
  74         self.collection = self.layer_collection.collection
  75
  76     def ensure_rig_object(self) -> tuple[bool, ArmatureObject]:
  77         """Check if the generated rig already exists, so we can
  78         regenerate in the same object. If not, create a new
  79         object to generate the rig in.
  80         """
  81         print("Fetch rig.")
  82         meta_data = self.metarig.data
  83
  84         target_rig = get_rigify_target_rig(meta_data)
  85         found = bool(target_rig)
  86
  87         if not found:
  88             rig_basename = get_rigify_rig_basename(meta_data)
  89
  90             if rig_basename:
  91                 rig_new_name = rig_basename
  92             elif "metarig" in self.metarig.name:
  93                 rig_new_name = self.metarig.name.replace("metarig", "rig")
  94             elif "META" in self.metarig.name:
  95                 rig_new_name = self.metarig.name.replace("META", "RIG")
  96             else:
  97                 rig_new_name = "RIG-" + self.metarig.name
  98
  99             arm = bpy.data.armatures.new(rig_new_name)
 100             target_rig = verify_armature_obj(bpy.data.objects.new(rig_new_name, arm))
 101             target_rig.display_type = 'WIRE'
 102
 103         # If the object is already added to the scene, switch to its collection
 104         if target_rig in list(self.context.scene.collection.all_objects):
 105             self.__switch_to_usable_collection(target_rig)
 106         else:
 107             # Otherwise, add to the selected collection or the metarig collection if unusable
 108             if (self.layer_collection not in self.usable_collections
 109                     or self.layer_collection == self.view_layer.layer_collection):
 110                 self.__switch_to_usable_collection(self.metarig, True)
 111
 112             self.collection.objects.link(target_rig)
 113
 114         # Configure and remember the object
 115         meta_data.rigify_target_rig = target_rig
 116         target_rig.data.pose_position = 'POSE'
 117
 118         return found, target_rig
 119
 120     def __unhide_rig_object(self, obj: bpy.types.Object):
 121         # Ensure the object is visible and selectable
 122         obj.hide_set(False, view_layer=self.view_layer)
 123         obj.hide_viewport = False
 124
 125         if not obj.visible_get(view_layer=self.view_layer):
 126             raise Exception('Could not generate: Target rig is not visible')
 127
 128         obj.select_set(True, view_layer=self.view_layer)
 129
 130         if not obj.select_get(view_layer=self.view_layer):
 131             raise Exception('Could not generate: Cannot select target rig')
 132
 133         if self.layer_collection not in self.usable_collections:
 134             raise Exception('Could not generate: Could not find a usable collection.')
 135
 136     def __save_rig_data(self, obj: ArmatureObject, obj_found: bool):
 137         if obj_found:
 138             self.saved_visible_layers = {coll.name: coll.is_visible for coll in obj.data.collections_all}
 139
 140             self.artifacts.generate_init_existing(obj)
 141
 142     def __find_legacy_collection(self) -> bpy.types.Collection:
 143         """For backwards comp, matching by name to find a legacy collection.
 144         (For before there was a Widget Collection PointerProperty)
 145         """
 146         widgets_group_name = WGT_GROUP_PREFIX + self.obj.name
 147         old_collection = bpy.data.collections.get(widgets_group_name)
 148
 149         if old_collection and old_collection.library:
 150             old_collection = None
 151
 152         if not old_collection:
 153             # Update the old 'Widgets' collection
 154             legacy_collection = bpy.data.collections.get('Widgets')
 155
 156             if legacy_collection and widgets_group_name in legacy_collection.objects\
 157                     and not legacy_collection.library:
 158                 legacy_collection.name = widgets_group_name
 159                 old_collection = legacy_collection
 160
 161         if old_collection:
 162             # Rename the collection
 163             old_collection.name = widgets_group_name
 164
 165         return old_collection
 166
 167     def ensure_widget_collection(self):
 168         # Create/find widget collection
 169         self.widget_collection = self.metarig.data.rigify_widgets_collection
 170         if not self.widget_collection:
 171             self.widget_collection = self.__find_legacy_collection()
 172         if not self.widget_collection:
 173             widgets_group_name = WGT_GROUP_PREFIX + self.obj.name.replace("RIG-", "")
 174             self.widget_collection = ensure_collection(
 175                 self.context, widgets_group_name, hidden=True)
 176
 177         self.metarig.data.rigify_widgets_collection = self.widget_collection
 178
 179         self.use_mirror_widgets = get_rigify_mirror_widgets(self.metarig.data)
 180
 181         # Build tables for existing widgets
 182         self.old_widget_table = {}
 183         self.new_widget_table = {}
 184         self.widget_mirror_mesh = {}
 185
 186         if get_rigify_force_widget_update(self.metarig.data):
 187             # Remove widgets if force update is set
 188             for obj in list(self.widget_collection.objects):
 189                 bpy.data.objects.remove(obj)
 190         elif self.obj.pose:
 191             # Find all widgets from the collection referenced by the old rig
 192             known_widgets = set(obj.name for obj in self.widget_collection.objects)
 193
 194             for bone in self.obj.pose.bones:
 195                 if bone.custom_shape and bone.custom_shape.name in known_widgets:
 196                     self.old_widget_table[bone.name] = bone.custom_shape
 197
 198             # Rename widgets in case the rig was renamed
 199             name_prefix = WGT_PREFIX + self.obj.name + "_"
 200
 201             for bone_name, widget in self.old_widget_table.items():
 202                 old_data_name = change_name_side(widget.name, get_name_side(widget.data.name))
 203
 204                 widget.name = name_prefix + bone_name
 205
 206                 # If the mesh name is the same as the object, rename it too
 207                 if widget.data.name == old_data_name:
 208                     widget.data.name = change_name_side(
 209                         widget.name, get_name_side(widget.data.name))
 210
 211             # Find meshes for mirroring
 212             if self.use_mirror_widgets:
 213                 for bone_name, widget in self.old_widget_table.items():
 214                     mid_name = change_name_side(bone_name, Side.MIDDLE)
 215                     if bone_name != mid_name:
 216                         assert isinstance(widget.data, bpy.types.Mesh)
 217                         self.widget_mirror_mesh[mid_name] = widget.data
 218
 219     def ensure_root_bone_collection(self):
 220         collections = self.metarig.data.collections_all
 221
 222         validate_collection_references(self.metarig)
 223
 224         coll = collections.get(ROOT_COLLECTION)
 225
 226         if not coll:
 227             coll = self.metarig.data.collections.new(ROOT_COLLECTION)
 228
 229         if coll.rigify_ui_row <= 0:
 230             coll.rigify_ui_row = 2 + choose_next_uid(collections, 'rigify_ui_row', min_value=1)
 231
 232     def __duplicate_rig(self):
 233         obj = self.obj
 234         metarig = self.metarig
 235         context = self.context
 236
 237         # Remove all bones from the generated rig armature.
 238         bpy.ops.object.mode_set(mode='EDIT')
 239         for bone in obj.data.edit_bones:
 240             obj.data.edit_bones.remove(bone)
 241         bpy.ops.object.mode_set(mode='OBJECT')
 242
 243         # Remove all bone collections from the target armature.
 244         for coll in list(obj.data.collections_all):
 245             obj.data.collections.remove(coll)
 246
 247         # Select and duplicate metarig
 248         select_object(context, metarig, deselect_all=True)
 249
 250         bpy.ops.object.duplicate()
 251
 252         # Rename org bones in the temporary object
 253         temp_obj = verify_armature_obj(context.view_layer.objects.active)
 254
 255         assert temp_obj and temp_obj != metarig
 256
 257         self.__freeze_driver_vars(temp_obj)
 258         self.__rename_org_bones(temp_obj)
 259
 260         # Select the target rig and join
 261         select_object(context, obj)
 262
 263         saved_matrix = obj.matrix_world.copy()
 264         obj.matrix_world = metarig.matrix_world
 265
 266         bpy.ops.object.join()
 267
 268         obj.matrix_world = saved_matrix
 269
 270         # Select the generated rig
 271         select_object(context, obj, deselect_all=True)
 272
 273         # Clean up animation data
 274         if obj.animation_data:
 275             obj.animation_data.action = None
 276
 277             for track in obj.animation_data.nla_tracks:
 278                 obj.animation_data.nla_tracks.remove(track)
 279
 280     @staticmethod
 281     def __freeze_driver_vars(obj: bpy.types.Object):
 282         if obj.animation_data:
 283             # Freeze drivers referring to custom properties
 284             for d in obj.animation_data.drivers:
 285                 for var in d.driver.variables:
 286                     for tar in var.targets:
 287                         # If a custom property
 288                         if var.type == 'SINGLE_PROP' \
 289                                 and re.match(r'^pose.bones\["[^"\]]*"]\["[^"\]]*"]$',
 290                                              tar.data_path):
 291                             tar.data_path = "RIGIFY-" + tar.data_path
 292
 293     def __rename_org_bones(self, obj: ArmatureObject):
 294         # Make a list of the original bones, so we can keep track of them.
 295         original_bones = [bone.name for bone in obj.data.bones]
 296
 297         # Add the ORG_PREFIX to the original bones.
 298         for i in range(0, len(original_bones)):
 299             bone = obj.pose.bones[original_bones[i]]
 300
 301             # Preserve the root bone as is if present
 302             if bone.name == ROOT_NAME:
 303                 if bone.parent:
 304                     raise MetarigError('Root bone must have no parent')
 305                 if get_rigify_type(bone) not in ('', 'basic.raw_copy'):
 306                     raise MetarigError('Root bone must have no rig, or use basic.raw_copy')
 307                 continue
 308
 309             # This rig type is special in that it preserves the name of the bone.
 310             if get_rigify_type(bone) != 'basic.raw_copy':
 311                 bone.name = make_original_name(original_bones[i])
 312                 original_bones[i] = bone.name
 313
 314         self.original_bones = original_bones
 315
 316     def __create_root_bone(self):
 317         obj = self.obj
 318         metarig = self.metarig
 319
 320         if ROOT_NAME in obj.data.bones:
 321             # Use the existing root bone
 322             root_bone = ROOT_NAME
 323         else:
 324             # Create the root bone.
 325             root_bone = new_bone(obj, ROOT_NAME)
 326             spread = get_xy_spread(metarig.data.bones) or metarig.data.bones[0].length
 327             spread = float('%.3g' % spread)
 328             scale = spread/0.589
 329             obj.data.edit_bones[root_bone].head = (0, 0, 0)
 330             obj.data.edit_bones[root_bone].tail = (0, scale, 0)
 331             obj.data.edit_bones[root_bone].roll = 0
 332
 333         self.root_bone = root_bone
 334         self.bone_owners[root_bone] = None
 335         self.noparent_bones.add(root_bone)
 336
 337     def __parent_bones_to_root(self):
 338         eb = self.obj.data.edit_bones
 339
 340         # Parent loose bones to root
 341         for bone in eb:
 342             if bone.name in self.noparent_bones:
 343                 continue
 344             elif bone.parent is None:
 345                 bone.use_connect = False
 346                 bone.parent = eb[self.root_bone]
 347
 348     def __lock_transforms(self):
 349         # Lock transforms on all non-control bones
 350         r = re.compile("[A-Z][A-Z][A-Z]-")
 351         for pb in self.obj.pose.bones:
 352             if r.match(pb.name):
 353                 pb.lock_location = (True, True, True)
 354                 pb.lock_rotation = (True, True, True)
 355                 pb.lock_rotation_w = True
 356                 pb.lock_scale = (True, True, True)
 357
 358     def ensure_bone_collection(self, name):
 359         coll = self.obj.data.collections_all.get(name)
 360
 361         if not coll:
 362             coll = self.obj.data.collections.new(name)
 363
 364         return coll
 365
 366     def __assign_layers(self):
 367         pose_bones = self.obj.pose.bones
 368
 369         root_coll = self.ensure_bone_collection(ROOT_COLLECTION)
 370         org_coll = self.ensure_bone_collection(ORG_COLLECTION)
 371         mch_coll = self.ensure_bone_collection(MCH_COLLECTION)
 372         def_coll = self.ensure_bone_collection(DEF_COLLECTION)
 373
 374         set_bone_layers(pose_bones[self.root_bone].bone, [root_coll])
 375
 376         # Every bone that has a name starting with "DEF-" make deforming.  All the
 377         # others make non-deforming.
 378         for pbone in pose_bones:
 379             bone = pbone.bone
 380             name = bone.name
 381             layers = None
 382
 383             bone.use_deform = name.startswith(DEF_PREFIX)
 384
 385             # Move all the original bones to their layer.
 386             if name.startswith(ORG_PREFIX):
 387                 layers = [org_coll]
 388             # Move all the bones with names starting with "MCH-" to their layer.
 389             elif name.startswith(MCH_PREFIX):
 390                 layers = [mch_coll]
 391             # Move all the bones with names starting with "DEF-" to their layer.
 392             elif name.startswith(DEF_PREFIX):
 393                 layers = [def_coll]
 394
 395             if layers is not None:
 396                 set_bone_layers(bone, layers)
 397
 398                 # Remove custom shapes from non-control bones
 399                 pbone.custom_shape = None
 400
 401             bone.bbone_x = bone.bbone_z = bone.length * 0.05
 402
 403     def __restore_driver_vars(self):
 404         obj = self.obj
 405
 406         # Alter marked driver targets
 407         if obj.animation_data:
 408             for d in obj.animation_data.drivers:
 409                 for v in d.driver.variables:
 410                     for tar in v.targets:
 411                         if tar.data_path.startswith("RIGIFY-"):
 412                             temp, bone, prop = tuple(
 413                                 [x.strip('"]') for x in tar.data_path.split('["')])
 414                             if bone in obj.data.bones and prop in obj.pose.bones[bone].keys():
 415                                 tar.data_path = tar.data_path[7:]
 416                             else:
 417                                 org_name = make_original_name(bone)
 418                                 org_name = self.org_rename_table.get(org_name, org_name)
 419                                 tar.data_path = 'pose.bones["%s"]["%s"]' % (org_name, prop)
 420
 421     def __assign_widgets(self):
 422         obj_table = {obj.name: obj for obj in self.scene.objects}
 423
 424         # Assign shapes to bones
 425         # Object's with name WGT-<bone_name> get used as that bone's shape.
 426         for bone in self.obj.pose.bones:
 427             # First check the table built by create_widget
 428             if bone.name in self.new_widget_table:
 429                 bone.custom_shape = self.new_widget_table[bone.name]
 430                 continue
 431
 432             # Object names are limited to 63 characters... arg
 433             wgt_name = (WGT_PREFIX + self.obj.name + '_' + bone.name)[:63]
 434
 435             if wgt_name in obj_table:
 436                 bone.custom_shape = obj_table[wgt_name]
 437
 438     def __compute_visible_layers(self):
 439         has_ui_buttons = set().union(*[
 440             {p.name for p in flatten_parents(coll)}
 441             for coll in self.obj.data.collections_all
 442             if coll.rigify_ui_row > 0
 443         ])
 444
 445         # Hide all layers without UI buttons
 446         for coll in self.obj.data.collections_all:
 447             user_visible = self.saved_visible_layers.get(coll.name, coll.is_visible)
 448             coll.is_visible = user_visible and coll.name in has_ui_buttons
 449
 450     def generate(self):
 451         context = self.context
 452         metarig = self.metarig
 453         view_layer = self.view_layer
 454         t = Timer()
 455
 456         self.usable_collections = list_layer_collections(
 457             view_layer.layer_collection, selectable=True)
 458
 459         bpy.ops.object.mode_set(mode='OBJECT')
 460
 461         ###########################################
 462         # Create/find the rig object and set it up
 463         obj_found, obj = self.ensure_rig_object()
 464
 465         self.obj = obj
 466         self.__unhide_rig_object(obj)
 467
 468         # Collect data from the existing rig
 469         self.artifacts = ArtifactManager(self)
 470
 471         self.__save_rig_data(obj, obj_found)
 472
 473         # Select the chosen working collection in case it changed
 474         self.view_layer.active_layer_collection = self.layer_collection
 475
 476         self.ensure_root_bone_collection()
 477
 478         # Get rid of anim data in case the rig already existed
 479         print("Clear rig animation data.")
 480
 481         obj.animation_data_clear()
 482         obj.data.animation_data_clear()
 483
 484         select_object(context, obj, deselect_all=True)
 485
 486         ###########################################
 487         # Create Widget Collection
 488         self.ensure_widget_collection()
 489
 490         t.tick("Create widgets collection: ")
 491
 492         ###########################################
 493         # Get parented objects to restore later
 494
 495         child_parent_bones = {}  # {object: bone}
 496
 497         for child in obj.children:
 498             child_parent_bones[child] = child.parent_bone
 499
 500         ###########################################
 501         # Copy bones from metarig to obj (adds ORG_PREFIX)
 502         self.__duplicate_rig()
 503
 504         obj.data.use_mirror_x = False
 505
 506         t.tick("Duplicate rig: ")
 507
 508         ###########################################
 509         # Put the rig_name in the armature custom properties
 510         obj.data["rig_id"] = self.rig_id
 511
 512         self.script = rig_ui_template.ScriptGenerator(self)
 513         self.action_layers = ActionLayerBuilder(self)
 514
 515         ###########################################
 516         bpy.ops.object.mode_set(mode='OBJECT')
 517
 518         self.instantiate_rig_tree()
 519
 520         t.tick("Instantiate rigs: ")
 521
 522         ###########################################
 523         bpy.ops.object.mode_set(mode='OBJECT')
 524
 525         self.invoke_initialize()
 526
 527         t.tick("Initialize rigs: ")
 528
 529         ###########################################
 530         bpy.ops.object.mode_set(mode='EDIT')
 531
 532         self.invoke_prepare_bones()
 533
 534         t.tick("Prepare bones: ")
 535
 536         ###########################################
 537         bpy.ops.object.mode_set(mode='OBJECT')
 538         bpy.ops.object.mode_set(mode='EDIT')
 539
 540         self.__create_root_bone()
 541
 542         self.invoke_generate_bones()
 543
 544         t.tick("Generate bones: ")
 545
 546         ###########################################
 547         bpy.ops.object.mode_set(mode='OBJECT')
 548         bpy.ops.object.mode_set(mode='EDIT')
 549
 550         self.invoke_parent_bones()
 551
 552         self.__parent_bones_to_root()
 553
 554         t.tick("Parent bones: ")
 555
 556         ###########################################
 557         bpy.ops.object.mode_set(mode='OBJECT')
 558
 559         self.invoke_configure_bones()
 560
 561         t.tick("Configure bones: ")
 562
 563         ###########################################
 564         bpy.ops.object.mode_set(mode='OBJECT')
 565
 566         self.invoke_preapply_bones()
 567
 568         t.tick("Preapply bones: ")
 569
 570         ###########################################
 571         bpy.ops.object.mode_set(mode='EDIT')
 572
 573         self.invoke_apply_bones()
 574
 575         t.tick("Apply bones: ")
 576
 577         ###########################################
 578         bpy.ops.object.mode_set(mode='OBJECT')
 579
 580         self.invoke_rig_bones()
 581
 582         t.tick("Rig bones: ")
 583
 584         ###########################################
 585         bpy.ops.object.mode_set(mode='OBJECT')
 586
 587         self.invoke_generate_widgets()
 588
 589         # Generate the default root widget last in case it's rigged with raw_copy
 590         create_root_widget(obj, self.root_bone)
 591
 592         t.tick("Generate widgets: ")
 593
 594         ###########################################
 595         bpy.ops.object.mode_set(mode='OBJECT')
 596
 597         self.__lock_transforms()
 598         self.__assign_layers()
 599         self.__compute_visible_layers()
 600         self.__restore_driver_vars()
 601
 602         t.tick("Assign layers: ")
 603
 604         ###########################################
 605         bpy.ops.object.mode_set(mode='OBJECT')
 606
 607         self.invoke_finalize()
 608
 609         t.tick("Finalize: ")
 610
 611         ###########################################
 612         bpy.ops.object.mode_set(mode='OBJECT')
 613
 614         self.__assign_widgets()
 615
 616         # Create Selection Sets
 617         create_selection_sets(obj, metarig)
 618
 619         # Create Bone Groups
 620         apply_bone_colors(obj, metarig, self.layer_group_priorities)
 621
 622         t.tick("The rest: ")
 623
 624         ###########################################
 625         # Restore state
 626
 627         bpy.ops.object.mode_set(mode='OBJECT')
 628         obj.data.pose_position = 'POSE'
 629
 630         # Restore parent to bones
 631         for child, sub_parent in child_parent_bones.items():
 632             if sub_parent in obj.pose.bones:
 633                 mat = child.matrix_world.copy()
 634                 child.parent_bone = sub_parent
 635                 child.matrix_world = mat
 636
 637         # Clear any transient errors in drivers
 638         refresh_all_drivers()
 639
 640         ###########################################
 641         # Execute the finalize script
 642
 643         finalize_script = get_rigify_finalize_script(metarig.data)
 644
 645         if finalize_script:
 646             bpy.ops.object.mode_set(mode='OBJECT')
 647             exec(finalize_script.as_string(), {})
 648             bpy.ops.object.mode_set(mode='OBJECT')
 649
 650         obj.data.collections.active_index = 0
 651
 652         self.artifacts.generate_cleanup()
 653
 654         ###########################################
 655         # Restore active collection
 656         view_layer.active_layer_collection = self.layer_collection
 657
 658
 659 def generate_rig(context, metarig):
 660     """ Generates a rig from a metarig.
 661
 662     """
 663     # Initial configuration
 664     rest_backup = metarig.data.pose_position
 665     metarig.data.pose_position = 'REST'
 666
 667     try:
 668         generator = Generator(context, metarig)
 669
 670         base_generate.BaseGenerator.instance = generator
 671
 672         generator.generate()
 673
 674         metarig.data.pose_position = rest_backup
 675
 676     except Exception as e:
 677         # Cleanup if something goes wrong
 678         print("Rigify: failed to generate rig.")
 679
 680         bpy.ops.object.mode_set(mode='OBJECT')
 681         metarig.data.pose_position = rest_backup
 682
 683         # Continue the exception
 684         raise e
 685
 686     finally:
 687         base_generate.BaseGenerator.instance = None
 688
 689
 690 def create_selection_set_for_rig_layer(rig: ArmatureObject, set_name: str, coll: bpy.types.BoneCollection) -> None:
 691     """Create a single selection set on a rig.
 692
 693     The set will contain all bones on the rig layer with the given index.
 694     """
 695     sel_set = rig.selection_sets.add()  # noqa
 696     sel_set.name = set_name
 697
 698     for b in rig.pose.bones:
 699         if coll not in b.bone.collections or b.name in sel_set.bone_ids:
 700             continue
 701
 702         bone_id = sel_set.bone_ids.add()
 703         bone_id.name = b.name
 704
 705
 706 def create_selection_sets(obj: ArmatureObject, _metarig: ArmatureObject):
 707     """Create selection sets if the Selection Sets addon is enabled.
 708
 709     Whether a selection set for a rig layer is created is controlled in the
 710     Rigify Layer Names panel.
 711     """
 712     # Check if selection sets addon is installed
 713     if 'bone_selection_groups' not in bpy.context.preferences.addons \
 714             and 'bone_selection_sets' not in bpy.context.preferences.addons:
 715         return
 716
 717     obj.selection_sets.clear()  # noqa
 718
 719     for coll in obj.data.collections_all:
 720         if not coll.rigify_sel_set:
 721             continue
 722
 723         create_selection_set_for_rig_layer(obj, coll.name, coll)
 724
 725
 726 def apply_bone_colors(obj, metarig, priorities: Optional[dict[str, dict[str, float]]] = None):
 727     bpy.ops.object.mode_set(mode='OBJECT')
 728     pb = obj.pose.bones
 729
 730     color_sets = get_rigify_colors(metarig.data)
 731     color_map = {i+1: cset for i, cset in enumerate(color_sets)}
 732
 733     collection_table: dict[str, tuple[int, 'RigifyColorSet']] = {
 734         coll.name: (i, color_map[coll.rigify_color_set_id])
 735         for i, coll in enumerate(flatten_children(obj.data.collections))
 736         if coll.rigify_color_set_id in color_map
 737     }
 738
 739     priorities = priorities or {}
 740     dummy = {}
 741
 742     for b in pb:
 743         bone_priorities = priorities.get(b.name, dummy)
 744         cset_collections = [coll.name for coll in b.bone.collections if coll.name in collection_table]
 745         if cset_collections:
 746             best_name = max(
 747                 cset_collections,
 748                 key=lambda n: (bone_priorities.get(n, 0), -collection_table[n][0])
 749             )
 750             _, cset = collection_table[best_name]
 751             cset.apply(b.bone.color)
 752             cset.apply(b.color)
 753
 754
 755 def get_xy_spread(bones):
 756     x_max = 0
 757     y_max = 0
 758     for b in bones:
 759         x_max = max((x_max, abs(b.head[0]), abs(b.tail[0])))
 760         y_max = max((y_max, abs(b.head[1]), abs(b.tail[1])))
 761
 762     return max((x_max, y_max))