bone_selection_sets.py

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