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