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