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, Antony Riakiotakis, Dan Eicher",
  22     "version": (2, 0, 1),
  23     "blender": (2, 75, 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
  72
  73 class POSE_PT_selection_sets(Panel):
  74     bl_label = "Selection Sets"
  75     bl_space_type = 'PROPERTIES'
  76     bl_region_type = 'WINDOW'
  77     bl_context = "data"
  78
  79     @classmethod
  80     def poll(cls, context):
  81         return (context.object and
  82                 context.object.type == 'ARMATURE' and
  83                 context.object.pose)
  84
  85     def draw(self, context):
  86         layout = self.layout
  87
  88         arm = context.object
  89
  90         row = layout.row()
  91         row.enabled = (context.mode == 'POSE')
  92
  93         # UI list
  94         rows = 4 if len(arm.selection_sets) > 0 else 1
  95         row.template_list(
  96             "POSE_UL_selection_set", "",  # type and unique id
  97             arm, "selection_sets",  # pointer to the CollectionProperty
  98             arm, "active_selection_set",  # pointer to the active identifier
  99             rows=rows
 100         )
 101
 102         # add/remove/specials UI list Menu
 103         col = row.column(align=True)
 104         col.operator("pose.selection_set_add", icon='ZOOMIN', text="")
 105         col.operator("pose.selection_set_remove", icon='ZOOMOUT', text="")
 106         col.menu("POSE_MT_selection_sets_specials", icon='DOWNARROW_HLT', text="")
 107
 108         # move up/down arrows
 109         if len(arm.selection_sets) > 0:
 110             col.separator()
 111             col.operator("pose.selection_set_move", icon='TRIA_UP', text="").direction = 'UP'
 112             col.operator("pose.selection_set_move", icon='TRIA_DOWN', text="").direction = 'DOWN'
 113
 114         # buttons
 115         row = layout.row()
 116
 117         sub = row.row(align=True)
 118         sub.operator("pose.selection_set_assign", text="Assign")
 119         sub.operator("pose.selection_set_unassign", text="Remove")
 120
 121         sub = row.row(align=True)
 122         sub.operator("pose.selection_set_select", text="Select")
 123         sub.operator("pose.selection_set_deselect", text="Deselect")
 124
 125
 126 class POSE_UL_selection_set(UIList):
 127     def draw_item(self, context, layout, data, set, icon, active_data, active_propname, index):
 128         layout.prop(set, "name", text="", icon='GROUP_BONE', emboss=False)
 129
 130
 131 class POSE_MT_create_new_selection_set(Menu):
 132     bl_idname = "POSE_MT_selection_set_create"
 133     bl_label = "Choose Selection Set"
 134
 135     def draw(self, context):
 136         layout = self.layout
 137         layout.operator("pose.selection_set_add_and_assign",
 138             text="New Selection Set")
 139
 140
 141 # Operators ###################################################################
 142
 143 class PluginOperator(Operator):
 144     @classmethod
 145     def poll(self, context):
 146         return (context.object and
 147                 context.object.type == 'ARMATURE' and
 148                 context.mode == 'POSE')
 149
 150
 151 class NeedSelSetPluginOperator(PluginOperator):
 152     @classmethod
 153     def poll(self, context):
 154         if super().poll(context):
 155             arm = context.object
 156             return (arm.active_selection_set < len(arm.selection_sets) and
 157                     arm.active_selection_set >= 0)
 158         return False
 159
 160
 161 class POSE_OT_selection_set_delete_all(PluginOperator):
 162     bl_idname = "pose.selection_set_delete_all"
 163     bl_label = "Delete All Sets"
 164     bl_description = "Deletes All Selection Sets"
 165     bl_options = {'UNDO', 'REGISTER'}
 166
 167     def execute(self, context):
 168         arm = context.object
 169         arm.selection_sets.clear()
 170         return {'FINISHED'}
 171
 172
 173 class POSE_OT_selection_set_remove_bones(PluginOperator):
 174     bl_idname = "pose.selection_set_remove_bones"
 175     bl_label = "Remove Bones from Sets"
 176     bl_description = "Removes the Active Bones from All Sets"
 177     bl_options = {'UNDO', 'REGISTER'}
 178
 179     def execute(self, context):
 180         arm = context.object
 181
 182         # iterate only the selected bones in current pose that are not hidden
 183         for bone in context.selected_pose_bones:
 184             for selset in arm.selection_sets:
 185                 if bone.name in selset.bone_ids:
 186                     idx = selset.bone_ids.find(bone.name)
 187                     selset.bone_ids.remove(idx)
 188
 189         return {'FINISHED'}
 190
 191
 192 class POSE_OT_selection_set_move(NeedSelSetPluginOperator):
 193     bl_idname = "pose.selection_set_move"
 194     bl_label = "Move Selection Set in List"
 195     bl_description = "Move the active Selection Set up/down the list of sets"
 196     bl_options = {'UNDO', 'REGISTER'}
 197
 198     direction = EnumProperty(
 199         name="Move Direction",
 200         description="Direction to move the active Selection Set: UP (default) or DOWN",
 201         items=[
 202             ('UP', "Up", "", -1),
 203             ('DOWN', "Down", "", 1),
 204         ],
 205         default='UP'
 206     )
 207
 208     @classmethod
 209     def poll(self, context):
 210         if super().poll(context):
 211             arm = context.object
 212             return len(arm.selection_sets) > 1
 213         return False
 214
 215     def execute(self, context):
 216         arm = context.object
 217
 218         active_idx = arm.active_selection_set
 219         new_idx = active_idx + (-1 if self.direction == 'UP' else 1)
 220
 221         if new_idx < 0 or new_idx >= len(arm.selection_sets):
 222             return {'FINISHED'}
 223
 224         arm.selection_sets.move(active_idx, new_idx)
 225         arm.active_selection_set = new_idx
 226
 227         return {'FINISHED'}
 228
 229
 230 class POSE_OT_selection_set_add(PluginOperator):
 231     bl_idname = "pose.selection_set_add"
 232     bl_label = "Create Selection Set"
 233     bl_description = "Creates a new empty Selection Set"
 234     bl_options = {'UNDO', 'REGISTER'}
 235
 236     def execute(self, context):
 237         arm = context.object
 238
 239         new_sel_set = arm.selection_sets.add()
 240
 241         # naming
 242         if "SelectionSet" not in arm.selection_sets:
 243             new_sel_set.name = "SelectionSet"
 244         else:
 245             sorted_sets = []
 246             for selset in arm.selection_sets:
 247                 if selset.name.startswith("SelectionSet."):
 248                     index = selset.name[13:]
 249                     if index.isdigit():
 250                         sorted_sets.append(index)
 251             sorted_sets = sorted(sorted_sets)
 252             min_index = 1
 253             for num in sorted_sets:
 254                 num = int(num)
 255                 if min_index < num:
 256                     break
 257                 min_index = num + 1
 258             new_sel_set.name = "SelectionSet.{:03d}".format(min_index)
 259
 260         # select newly created set
 261         arm.active_selection_set = len(arm.selection_sets) - 1
 262
 263         return {'FINISHED'}
 264
 265
 266 class POSE_OT_selection_set_remove(NeedSelSetPluginOperator):
 267     bl_idname = "pose.selection_set_remove"
 268     bl_label = "Delete Selection Set"
 269     bl_description = "Delete a Selection Set"
 270     bl_options = {'UNDO', 'REGISTER'}
 271
 272     def execute(self, context):
 273         arm = context.object
 274
 275         arm.selection_sets.remove(arm.active_selection_set)
 276
 277         # change currently active selection set
 278         numsets = len(arm.selection_sets)
 279         if (arm.active_selection_set > (numsets - 1) and numsets > 0):
 280             arm.active_selection_set = len(arm.selection_sets) - 1
 281
 282         return {'FINISHED'}
 283
 284
 285 class POSE_OT_selection_set_assign(PluginOperator):
 286     bl_idname = "pose.selection_set_assign"
 287     bl_label = "Add Bones to Selection Set"
 288     bl_description = "Add selected bones to Selection Set"
 289     bl_options = {'UNDO', 'REGISTER'}
 290
 291     def invoke(self, context, event):
 292         arm = context.object
 293
 294         if not (arm.active_selection_set < len(arm.selection_sets)):
 295             bpy.ops.wm.call_menu("INVOKE_DEFAULT",
 296                 name="POSE_MT_selection_set_create")
 297         else:
 298             bpy.ops.pose.selection_set_assign('EXEC_DEFAULT')
 299
 300         return {'FINISHED'}
 301
 302     def execute(self, context):
 303         arm = context.object
 304         act_sel_set = arm.selection_sets[arm.active_selection_set]
 305
 306         # iterate only the selected bones in current pose that are not hidden
 307         for bone in context.selected_pose_bones:
 308             if bone.name not in act_sel_set.bone_ids:
 309                 bone_id = act_sel_set.bone_ids.add()
 310                 bone_id.name = bone.name
 311
 312         return {'FINISHED'}
 313
 314
 315 class POSE_OT_selection_set_unassign(NeedSelSetPluginOperator):
 316     bl_idname = "pose.selection_set_unassign"
 317     bl_label = "Remove Bones from Selection Set"
 318     bl_description = "Remove selected bones from Selection Set"
 319     bl_options = {'UNDO', 'REGISTER'}
 320
 321     def execute(self, context):
 322         arm = context.object
 323         act_sel_set = arm.selection_sets[arm.active_selection_set]
 324
 325         # iterate only the selected bones in current pose that are not hidden
 326         for bone in context.selected_pose_bones:
 327             if bone.name in act_sel_set.bone_ids:
 328                 idx = act_sel_set.bone_ids.find(bone.name)
 329                 act_sel_set.bone_ids.remove(idx)
 330
 331         return {'FINISHED'}
 332
 333
 334 class POSE_OT_selection_set_select(NeedSelSetPluginOperator):
 335     bl_idname = "pose.selection_set_select"
 336     bl_label = "Select Selection Set"
 337     bl_description = "Add Selection Set bones to current selection"
 338     bl_options = {'UNDO', 'REGISTER'}
 339
 340     def execute(self, context):
 341         arm = context.object
 342         act_sel_set = arm.selection_sets[arm.active_selection_set]
 343
 344         for bone in context.visible_pose_bones:
 345             if bone.name in act_sel_set.bone_ids:
 346                 bone.bone.select = True
 347
 348         return {'FINISHED'}
 349
 350
 351 class POSE_OT_selection_set_deselect(NeedSelSetPluginOperator):
 352     bl_idname = "pose.selection_set_deselect"
 353     bl_label = "Deselect Selection Set"
 354     bl_description = "Remove Selection Set bones from current selection"
 355     bl_options = {'UNDO', 'REGISTER'}
 356
 357     def execute(self, context):
 358         arm = context.object
 359         act_sel_set = arm.selection_sets[arm.active_selection_set]
 360
 361         for bone in context.selected_pose_bones:
 362             if bone.name in act_sel_set.bone_ids:
 363                 bone.bone.select = False
 364
 365         return {'FINISHED'}
 366
 367
 368 class POSE_OT_selection_set_add_and_assign(PluginOperator):
 369     bl_idname = "pose.selection_set_add_and_assign"
 370     bl_label = "Create and Add Bones to Selection Set"
 371     bl_description = "Creates a new Selection Set with the currently selected bones"
 372     bl_options = {'UNDO', 'REGISTER'}
 373
 374     def execute(self, context):
 375         bpy.ops.pose.selection_set_add('EXEC_DEFAULT')
 376         bpy.ops.pose.selection_set_assign('EXEC_DEFAULT')
 377         return {'FINISHED'}
 378
 379
 380 # Registry ####################################################################
 381
 382 classes = (
 383     POSE_MT_create_new_selection_set,
 384     POSE_MT_selection_sets_specials,
 385     POSE_PT_selection_sets,
 386     POSE_UL_selection_set,
 387     SelectionEntry,
 388     SelectionSet,
 389     POSE_OT_selection_set_delete_all,
 390     POSE_OT_selection_set_remove_bones,
 391     POSE_OT_selection_set_move,
 392     POSE_OT_selection_set_add,
 393     POSE_OT_selection_set_remove,
 394     POSE_OT_selection_set_assign,
 395     POSE_OT_selection_set_unassign,
 396     POSE_OT_selection_set_select,
 397     POSE_OT_selection_set_deselect,
 398     POSE_OT_selection_set_add_and_assign,
 399 )
 400
 401
 402 def register():
 403     for cls in classes:
 404         bpy.utils.register_class(cls)
 405
 406     bpy.types.Object.selection_sets = CollectionProperty(
 407             type=SelectionSet,
 408             name="Selection Sets",
 409             description="List of groups of bones for easy selection"
 410             )
 411     bpy.types.Object.active_selection_set = IntProperty(
 412             name="Active Selection Set",
 413             description="Index of the currently active selection set",
 414             default=0
 415             )
 416
 417
 418 def unregister():
 419     for cls in classes:
 420         bpy.utils.unregister_class(cls)
 421
 422     del bpy.types.Object.selection_sets
 423     del bpy.types.Object.active_selection_set
 424
 425
 426 if __name__ == "__main__":
 427     register()