bone_selection_sets.py

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