bone_selection_sets.py

   1 # SPDX-FileCopyrightText: 2016-2022 Blender Foundation
   2 #
   3 # SPDX-License-Identifier: GPL-2.0-or-later
   4
   5 bl_info = {
   6     "name": "Bone Selection Sets",
   7     "author": "Inês Almeida, Sybren A. Stüvel, Antony Riakiotakis, Dan Eicher",
   8     "version": (2, 1, 1),
   9     "blender": (2, 80, 0),
  10     "location": "Properties > Object Data (Armature) > Selection Sets",
  11     "description": "List of Bone sets for easy selection while animating",
  12     "warning": "",
  13     "doc_url": "{BLENDER_MANUAL_URL}/addons/animation/bone_selection_sets.html",
  14     "category": "Animation",
  15 }
  16
  17 import bpy
  18 from bpy.types import (
  19     Operator,
  20     Menu,
  21     Panel,
  22     UIList,
  23     PropertyGroup,
  24 )
  25 from bpy.props import (
  26     StringProperty,
  27     IntProperty,
  28     EnumProperty,
  29     BoolProperty,
  30     CollectionProperty,
  31 )
  32
  33
  34 # Data Structure ##############################################################
  35
  36 # Note: bones are stored by name, this means that if the bone is renamed,
  37 # there can be problems. However, bone renaming is unlikely during animation.
  38 class SelectionEntry(PropertyGroup):
  39     name: StringProperty(name="Bone Name", override={'LIBRARY_OVERRIDABLE'})
  40
  41
  42 class SelectionSet(PropertyGroup):
  43     name: StringProperty(name="Set Name", override={'LIBRARY_OVERRIDABLE'})
  44     bone_ids: CollectionProperty(
  45         type=SelectionEntry,
  46         override={'LIBRARY_OVERRIDABLE', 'USE_INSERTION'}
  47     )
  48     is_selected: BoolProperty(name="Is Selected", override={'LIBRARY_OVERRIDABLE'})
  49
  50
  51 # UI Panel w/ UIList ##########################################################
  52
  53 class POSE_MT_selection_sets_context_menu(Menu):
  54     bl_label = "Selection Sets Specials"
  55
  56     def draw(self, context):
  57         layout = self.layout
  58
  59         layout.operator("pose.selection_set_delete_all", icon='X')
  60         layout.operator("pose.selection_set_remove_bones", icon='X')
  61         layout.operator("pose.selection_set_copy", icon='COPYDOWN')
  62         layout.operator("pose.selection_set_paste", icon='PASTEDOWN')
  63
  64
  65 class POSE_PT_selection_sets(Panel):
  66     bl_label = "Selection Sets"
  67     bl_space_type = 'PROPERTIES'
  68     bl_region_type = 'WINDOW'
  69     bl_context = "data"
  70     bl_options = {'DEFAULT_CLOSED'}
  71
  72     @classmethod
  73     def poll(cls, context):
  74         return (context.object and
  75                 context.object.type == 'ARMATURE' and
  76                 context.object.pose)
  77
  78     def draw(self, context):
  79         layout = self.layout
  80
  81         arm = context.object
  82
  83         row = layout.row()
  84         row.enabled = (context.mode == 'POSE')
  85
  86         # UI list
  87         rows = 4 if len(arm.selection_sets) > 0 else 1
  88         row.template_list(
  89             "POSE_UL_selection_set", "",  # type and unique id
  90             arm, "selection_sets",  # pointer to the CollectionProperty
  91             arm, "active_selection_set",  # pointer to the active identifier
  92             rows=rows
  93         )
  94
  95         # add/remove/specials UI list Menu
  96         col = row.column(align=True)
  97         col.operator("pose.selection_set_add", icon='ADD', text="")
  98         col.operator("pose.selection_set_remove", icon='REMOVE', text="")
  99         col.menu("POSE_MT_selection_sets_context_menu", icon='DOWNARROW_HLT', text="")
 100
 101         # move up/down arrows
 102         if len(arm.selection_sets) > 0:
 103             col.separator()
 104             col.operator("pose.selection_set_move", icon='TRIA_UP', text="").direction = 'UP'
 105             col.operator("pose.selection_set_move", icon='TRIA_DOWN', text="").direction = 'DOWN'
 106
 107         # buttons
 108         row = layout.row()
 109
 110         sub = row.row(align=True)
 111         sub.operator("pose.selection_set_assign", text="Assign")
 112         sub.operator("pose.selection_set_unassign", text="Remove")
 113
 114         sub = row.row(align=True)
 115         sub.operator("pose.selection_set_select", text="Select")
 116         sub.operator("pose.selection_set_deselect", text="Deselect")
 117
 118
 119 class POSE_UL_selection_set(UIList):
 120     def draw_item(self, context, layout, data, item, icon, active_data, active_propname, index):
 121         sel_set = item
 122         layout.prop(item, "name", text="", icon='GROUP_BONE', emboss=False)
 123         if self.layout_type in ('DEFAULT', 'COMPACT'):
 124             layout.prop(item, "is_selected", text="")
 125
 126
 127 class POSE_MT_selection_set_create(Menu):
 128     bl_label = "Choose Selection Set"
 129
 130     def draw(self, context):
 131         layout = self.layout
 132         layout.operator("pose.selection_set_add_and_assign",
 133                         text="New Selection Set")
 134
 135
 136 class POSE_MT_selection_sets_select(Menu):
 137     bl_label = 'Select Selection Set'
 138
 139     @classmethod
 140     def poll(cls, context):
 141         return POSE_OT_selection_set_select.poll(context)
 142
 143     def draw(self, context):
 144         layout = self.layout
 145         layout.operator_context = 'EXEC_DEFAULT'
 146         for idx, sel_set in enumerate(context.object.selection_sets):
 147             props = layout.operator(POSE_OT_selection_set_select.bl_idname, text=sel_set.name)
 148             props.selection_set_index = idx
 149
 150
 151 # Operators ###################################################################
 152
 153 class PluginOperator(Operator):
 154     """Operator only available for objects of type armature in pose mode."""
 155     @classmethod
 156     def poll(cls, context):
 157         return (context.object and
 158                 context.object.type == 'ARMATURE' and
 159                 context.mode == 'POSE')
 160
 161
 162 class NeedSelSetPluginOperator(PluginOperator):
 163     """Operator only available if the armature has a selected selection set."""
 164     @classmethod
 165     def poll(cls, context):
 166         if not super().poll(context):
 167             return False
 168         arm = context.object
 169         return 0 <= arm.active_selection_set < len(arm.selection_sets)
 170
 171
 172 class POSE_OT_selection_set_delete_all(PluginOperator):
 173     bl_idname = "pose.selection_set_delete_all"
 174     bl_label = "Delete All Sets"
 175     bl_description = "Deletes All Selection Sets"
 176     bl_options = {'UNDO', 'REGISTER'}
 177
 178     def execute(self, context):
 179         arm = context.object
 180         arm.selection_sets.clear()
 181         return {'FINISHED'}
 182
 183
 184 class POSE_OT_selection_set_remove_bones(PluginOperator):
 185     bl_idname = "pose.selection_set_remove_bones"
 186     bl_label = "Remove Selected Bones from All Sets"
 187     bl_description = "Removes the Selected Bones from All Sets"
 188     bl_options = {'UNDO', 'REGISTER'}
 189
 190     def execute(self, context):
 191         arm = context.object
 192
 193         # iterate only the selected bones in current pose that are not hidden
 194         for bone in context.selected_pose_bones:
 195             for selset in arm.selection_sets:
 196                 if bone.name in selset.bone_ids:
 197                     idx = selset.bone_ids.find(bone.name)
 198                     selset.bone_ids.remove(idx)
 199
 200         return {'FINISHED'}
 201
 202
 203 class POSE_OT_selection_set_move(NeedSelSetPluginOperator):
 204     bl_idname = "pose.selection_set_move"
 205     bl_label = "Move Selection Set in List"
 206     bl_description = "Move the active Selection Set up/down the list of sets"
 207     bl_options = {'UNDO', 'REGISTER'}
 208
 209     direction: EnumProperty(
 210         name="Move Direction",
 211         description="Direction to move the active Selection Set: UP (default) or DOWN",
 212         items=[
 213             ('UP', "Up", "", -1),
 214             ('DOWN', "Down", "", 1),
 215         ],
 216         default='UP',
 217         options={'HIDDEN'},
 218     )
 219
 220     @classmethod
 221     def poll(cls, context):
 222         if not super().poll(context):
 223             return False
 224         arm = context.object
 225         return len(arm.selection_sets) > 1
 226
 227     def execute(self, context):
 228         arm = context.object
 229
 230         active_idx = arm.active_selection_set
 231         new_idx = active_idx + (-1 if self.direction == 'UP' else 1)
 232
 233         if new_idx < 0 or new_idx >= len(arm.selection_sets):
 234             return {'FINISHED'}
 235
 236         arm.selection_sets.move(active_idx, new_idx)
 237         arm.active_selection_set = new_idx
 238
 239         return {'FINISHED'}
 240
 241
 242 class POSE_OT_selection_set_add(PluginOperator):
 243     bl_idname = "pose.selection_set_add"
 244     bl_label = "Create Selection Set"
 245     bl_description = "Creates a new empty Selection Set"
 246     bl_options = {'UNDO', 'REGISTER'}
 247
 248     def execute(self, context):
 249         arm = context.object
 250         sel_sets = arm.selection_sets
 251         new_sel_set = sel_sets.add()
 252         new_sel_set.name = uniqify("SelectionSet", sel_sets.keys())
 253
 254         # select newly created set
 255         arm.active_selection_set = len(sel_sets) - 1
 256
 257         return {'FINISHED'}
 258
 259
 260 class POSE_OT_selection_set_remove(NeedSelSetPluginOperator):
 261     bl_idname = "pose.selection_set_remove"
 262     bl_label = "Delete Selection Set"
 263     bl_description = "Delete a Selection Set"
 264     bl_options = {'UNDO', 'REGISTER'}
 265
 266     def execute(self, context):
 267         arm = context.object
 268
 269         arm.selection_sets.remove(arm.active_selection_set)
 270
 271         # change currently active selection set
 272         numsets = len(arm.selection_sets)
 273         if (arm.active_selection_set > (numsets - 1) and numsets > 0):
 274             arm.active_selection_set = len(arm.selection_sets) - 1
 275
 276         return {'FINISHED'}
 277
 278
 279 class POSE_OT_selection_set_assign(PluginOperator):
 280     bl_idname = "pose.selection_set_assign"
 281     bl_label = "Add Bones to Selection Set"
 282     bl_description = "Add selected bones to Selection Set"
 283     bl_options = {'UNDO', 'REGISTER'}
 284
 285     def invoke(self, context, event):
 286         arm = context.object
 287
 288         if not (arm.active_selection_set < len(arm.selection_sets)):
 289             bpy.ops.wm.call_menu("INVOKE_DEFAULT",
 290                                  name="POSE_MT_selection_set_create")
 291         else:
 292             bpy.ops.pose.selection_set_assign('EXEC_DEFAULT')
 293
 294         return {'FINISHED'}
 295
 296     def execute(self, context):
 297         arm = context.object
 298         act_sel_set = arm.selection_sets[arm.active_selection_set]
 299
 300         # iterate only the selected bones in current pose that are not hidden
 301         for bone in context.selected_pose_bones:
 302             if bone.name not in act_sel_set.bone_ids:
 303                 bone_id = act_sel_set.bone_ids.add()
 304                 bone_id.name = bone.name
 305
 306         return {'FINISHED'}
 307
 308
 309 class POSE_OT_selection_set_unassign(NeedSelSetPluginOperator):
 310     bl_idname = "pose.selection_set_unassign"
 311     bl_label = "Remove Bones from Selection Set"
 312     bl_description = "Remove selected bones from Selection Set"
 313     bl_options = {'UNDO', 'REGISTER'}
 314
 315     def execute(self, context):
 316         arm = context.object
 317         act_sel_set = arm.selection_sets[arm.active_selection_set]
 318
 319         # iterate only the selected bones in current pose that are not hidden
 320         for bone in context.selected_pose_bones:
 321             if bone.name in act_sel_set.bone_ids:
 322                 idx = act_sel_set.bone_ids.find(bone.name)
 323                 act_sel_set.bone_ids.remove(idx)
 324
 325         return {'FINISHED'}
 326
 327
 328 class POSE_OT_selection_set_select(NeedSelSetPluginOperator):
 329     bl_idname = "pose.selection_set_select"
 330     bl_label = "Select Selection Set"
 331     bl_description = "Add Selection Set bones to current selection"
 332     bl_options = {'UNDO', 'REGISTER'}
 333
 334     selection_set_index: IntProperty(
 335         name='Selection Set Index',
 336         default=-1,
 337         description='Which Selection Set to select; -1 uses the active Selection Set',
 338         options={'HIDDEN'},
 339     )
 340
 341     def execute(self, context):
 342         arm = context.object
 343
 344         if self.selection_set_index == -1:
 345             idx = arm.active_selection_set
 346         else:
 347             idx = self.selection_set_index
 348         sel_set = arm.selection_sets[idx]
 349
 350         for bone in context.visible_pose_bones:
 351             if bone.name in sel_set.bone_ids:
 352                 bone.bone.select = True
 353
 354         return {'FINISHED'}
 355
 356
 357 class POSE_OT_selection_set_deselect(NeedSelSetPluginOperator):
 358     bl_idname = "pose.selection_set_deselect"
 359     bl_label = "Deselect Selection Set"
 360     bl_description = "Remove Selection Set bones from current selection"
 361     bl_options = {'UNDO', 'REGISTER'}
 362
 363     def execute(self, context):
 364         arm = context.object
 365         act_sel_set = arm.selection_sets[arm.active_selection_set]
 366
 367         for bone in context.selected_pose_bones:
 368             if bone.name in act_sel_set.bone_ids:
 369                 bone.bone.select = False
 370
 371         return {'FINISHED'}
 372
 373
 374 class POSE_OT_selection_set_add_and_assign(PluginOperator):
 375     bl_idname = "pose.selection_set_add_and_assign"
 376     bl_label = "Create and Add Bones to Selection Set"
 377     bl_description = "Creates a new Selection Set with the currently selected bones"
 378     bl_options = {'UNDO', 'REGISTER'}
 379
 380     def execute(self, context):
 381         bpy.ops.pose.selection_set_add('EXEC_DEFAULT')
 382         bpy.ops.pose.selection_set_assign('EXEC_DEFAULT')
 383         return {'FINISHED'}
 384
 385
 386 class POSE_OT_selection_set_copy(NeedSelSetPluginOperator):
 387     bl_idname = "pose.selection_set_copy"
 388     bl_label = "Copy Selection Set(s)"
 389     bl_description = "Copies the selected Selection Set(s) to the clipboard"
 390     bl_options = {'UNDO', 'REGISTER'}
 391
 392     def execute(self, context):
 393         context.window_manager.clipboard = to_json(context)
 394         self.report({'INFO'}, 'Copied Selection Set(s) to Clipboard')
 395         return {'FINISHED'}
 396
 397
 398 class POSE_OT_selection_set_paste(PluginOperator):
 399     bl_idname = "pose.selection_set_paste"
 400     bl_label = "Paste Selection Set(s)"
 401     bl_description = "Adds new Selection Set(s) from the Clipboard"
 402     bl_options = {'UNDO', 'REGISTER'}
 403
 404     def execute(self, context):
 405         import json
 406
 407         try:
 408             from_json(context, context.window_manager.clipboard)
 409         except (json.JSONDecodeError, KeyError):
 410             self.report({'ERROR'}, 'The clipboard does not contain a Selection Set')
 411         else:
 412             # Select the pasted Selection Set.
 413             context.object.active_selection_set = len(context.object.selection_sets) - 1
 414
 415         return {'FINISHED'}
 416
 417
 418 # Helper Functions ############################################################
 419
 420 def menu_func_select_selection_set(self, context):
 421     self.layout.menu('POSE_MT_selection_sets_select', text="Bone Selection Set")
 422
 423
 424 def to_json(context) -> str:
 425     """Convert the selected Selection Sets of the current rig to JSON.
 426
 427     Selected Sets are the active_selection_set determined by the UIList
 428     plus any with the is_selected checkbox on."""
 429     import json
 430
 431     arm = context.object
 432     active_idx = arm.active_selection_set
 433
 434     json_obj = {}
 435     for idx, sel_set in enumerate(context.object.selection_sets):
 436         if idx == active_idx or sel_set.is_selected:
 437             bones = [bone_id.name for bone_id in sel_set.bone_ids]
 438             json_obj[sel_set.name] = bones
 439
 440     return json.dumps(json_obj)
 441
 442
 443 def from_json(context, as_json: str):
 444     """Add the selection sets (one or more) from JSON to the current rig."""
 445     import json
 446
 447     json_obj = json.loads(as_json)
 448     arm_sel_sets = context.object.selection_sets
 449
 450     for name, bones in json_obj.items():
 451         new_sel_set = arm_sel_sets.add()
 452         new_sel_set.name = uniqify(name, arm_sel_sets.keys())
 453         for bone_name in bones:
 454             bone_id = new_sel_set.bone_ids.add()
 455             bone_id.name = bone_name
 456
 457
 458 def uniqify(name: str, other_names: list) -> str:
 459     """Return a unique name with .xxx suffix if necessary.
 460
 461     Example usage:
 462
 463     >>> uniqify('hey', ['there'])
 464     'hey'
 465     >>> uniqify('hey', ['hey.001', 'hey.005'])
 466     'hey'
 467     >>> uniqify('hey', ['hey', 'hey.001', 'hey.005'])
 468     'hey.002'
 469     >>> uniqify('hey', ['hey', 'hey.005', 'hey.001'])
 470     'hey.002'
 471     >>> uniqify('hey', ['hey', 'hey.005', 'hey.001', 'hey.left'])
 472     'hey.002'
 473     >>> uniqify('hey', ['hey', 'hey.001', 'hey.002'])
 474     'hey.003'
 475
 476     It also works with a dict_keys object:
 477     >>> uniqify('hey', {'hey': 1, 'hey.005': 1, 'hey.001': 1}.keys())
 478     'hey.002'
 479     """
 480
 481     if name not in other_names:
 482         return name
 483
 484     # Construct the list of numbers already in use.
 485     offset = len(name) + 1
 486     others = (n[offset:] for n in other_names
 487               if n.startswith(name + '.'))
 488     numbers = sorted(int(suffix) for suffix in others
 489                      if suffix.isdigit())
 490
 491     # Find the first unused number.
 492     min_index = 1
 493     for num in numbers:
 494         if min_index < num:
 495             break
 496         min_index = num + 1
 497     return "{}.{:03d}".format(name, min_index)
 498
 499
 500 # Registry ####################################################################
 501
 502 classes = (
 503     POSE_MT_selection_set_create,
 504     POSE_MT_selection_sets_context_menu,
 505     POSE_MT_selection_sets_select,
 506     POSE_PT_selection_sets,
 507     POSE_UL_selection_set,
 508     SelectionEntry,
 509     SelectionSet,
 510     POSE_OT_selection_set_delete_all,
 511     POSE_OT_selection_set_remove_bones,
 512     POSE_OT_selection_set_move,
 513     POSE_OT_selection_set_add,
 514     POSE_OT_selection_set_remove,
 515     POSE_OT_selection_set_assign,
 516     POSE_OT_selection_set_unassign,
 517     POSE_OT_selection_set_select,
 518     POSE_OT_selection_set_deselect,
 519     POSE_OT_selection_set_add_and_assign,
 520     POSE_OT_selection_set_copy,
 521     POSE_OT_selection_set_paste,
 522 )
 523
 524
 525 # Store keymaps here to access after registration.
 526 addon_keymaps = []
 527
 528
 529 def register():
 530     for cls in classes:
 531         bpy.utils.register_class(cls)
 532
 533     # Add properties.
 534     bpy.types.Object.selection_sets = CollectionProperty(
 535         type=SelectionSet,
 536         name="Selection Sets",
 537         description="List of groups of bones for easy selection",
 538         override={'LIBRARY_OVERRIDABLE', 'USE_INSERTION'}
 539     )
 540     bpy.types.Object.active_selection_set = IntProperty(
 541         name="Active Selection Set",
 542         description="Index of the currently active selection set",
 543         default=0,
 544         override={'LIBRARY_OVERRIDABLE'}
 545     )
 546
 547     # Add shortcuts to the keymap.
 548     wm = bpy.context.window_manager
 549     if wm.keyconfigs.addon is not None:
 550         # wm.keyconfigs.addon is None when Blender is running in the background.
 551         km = wm.keyconfigs.addon.keymaps.new(name='Pose')
 552         kmi = km.keymap_items.new('wm.call_menu', 'W', 'PRESS', alt=True, shift=True)
 553         kmi.properties.name = 'POSE_MT_selection_sets_select'
 554         addon_keymaps.append((km, kmi))
 555
 556     # Add entries to menus.
 557     bpy.types.VIEW3D_MT_select_pose.append(menu_func_select_selection_set)
 558
 559
 560 def unregister():
 561     for cls in classes:
 562         bpy.utils.unregister_class(cls)
 563
 564     # Clear properties.
 565     del bpy.types.Object.selection_sets
 566     del bpy.types.Object.active_selection_set
 567
 568     # Clear shortcuts from the keymap.
 569     for km, kmi in addon_keymaps:
 570         km.keymap_items.remove(kmi)
 571     addon_keymaps.clear()
 572
 573     # Clear entries from menus.
 574     bpy.types.VIEW3D_MT_select_pose.remove(menu_func_select_selection_set)
 575
 576
 577
 578 if __name__ == "__main__":
 579     import doctest
 580
 581     doctest.testmod()
 582     register()