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": "Dan Eicher, Antony Riakiotakis, Inês Almeida",
  22     "version": (2, 0, 0),
  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": "http://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
  41 from bpy.props import (
  42     StringProperty,
  43     IntProperty,
  44     EnumProperty,
  45     CollectionProperty,
  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         # TODO
  70         #layout.operator("pose.selection_sets_sort", icon='SORTALPHA', text="Sort by Name").sort_type = 'NAME'
  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
  82             and context.object.type == 'ARMATURE'
  83             and context.object.pose)
  84
  85     def draw(self, context):
  86         layout = self.layout
  87
  88         ob = context.object
  89         arm = context.object
  90
  91         row = layout.row()
  92         row.enabled = (context.mode == 'POSE')
  93
  94         # UI list
  95         rows = 4  if len(arm.selection_sets) > 0 else 1
  96         row.template_list(
  97             "POSE_UL_selection_set", "", # type and unique id
  98             arm, "selection_sets", # pointer to the CollectionProperty
  99             arm, "active_selection_set", # pointer to the active identifier
 100             rows=rows
 101         )
 102
 103         # add/remove/specials UI list Menu
 104         col = row.column(align=True)
 105         col.operator("pose.selection_set_add", icon='ZOOMIN', text="")
 106         col.operator("pose.selection_set_remove", icon='ZOOMOUT', text="")
 107         # TODO specials like sorting
 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_create_new_selection_set(Menu):
 134     bl_idname = "pose.selection_set_create_new_popup"
 135     bl_label = "Choose Selection Set"
 136
 137     def draw(self, context):
 138         layout = self.layout
 139         layout.operator("pose.selection_set_add_and_assign",
 140             text="New Selection Set")
 141
 142
 143 # Operators ###################################################################
 144
 145 class PluginOperator(Operator):
 146     @classmethod
 147     def poll(self, context):
 148         return (context.object and
 149                 context.object.type == 'ARMATURE' and
 150                 context.mode == 'POSE')
 151
 152 class NeedSelSetPluginOperator(PluginOperator):
 153     @classmethod
 154     def poll(self, context):
 155         if super().poll(context):
 156             arm =  context.object
 157             return (arm.active_selection_set < len(arm.selection_sets)
 158                 and arm.active_selection_set >= 0)
 159         return False
 160
 161
 162 class POSE_OT_selection_set_move(NeedSelSetPluginOperator):
 163     bl_idname = "pose.selection_set_move"
 164     bl_label = "Move Selection Set in List"
 165     bl_description = "Move the active Selection Set up/down the list of sets"
 166     bl_options = {'UNDO', 'REGISTER'}
 167
 168     direction = EnumProperty(
 169         name="Move Direction",
 170         description="Direction to move the active Selection Set: UP (default) or DOWN",
 171         items=[
 172             ('UP', "Up", "", -1),
 173             ('DOWN', "Down", "", 1),
 174         ],
 175         default='UP'
 176     )
 177
 178     @classmethod
 179     def poll(self, context):
 180         if super().poll(context):
 181             arm =  context.object
 182             return len(arm.selection_sets) > 1
 183         return False
 184
 185     def execute(self, context):
 186         arm = context.object
 187
 188         active_idx = arm.active_selection_set
 189         new_idx = active_idx + (-1 if self.direction == 'UP' else 1)
 190
 191         if new_idx < 0 or new_idx >= len(arm.selection_sets):
 192             return {'FINISHED'}
 193
 194         arm.selection_sets.move(active_idx, new_idx)
 195         arm.active_selection_set = new_idx
 196
 197         return {'FINISHED'}
 198
 199
 200 class POSE_OT_selection_set_add(PluginOperator):
 201     bl_idname = "pose.selection_set_add"
 202     bl_label = "Create Selection Set"
 203     bl_description = "Creates a new empty Selection Set"
 204     bl_options = {'UNDO', 'REGISTER'}
 205
 206
 207     def execute(self, context):
 208         arm = context.object
 209
 210         new_sel_set = arm.selection_sets.add()
 211
 212         # naming
 213         if "SelectionSet" not in arm.selection_sets:
 214             new_sel_set.name  = "SelectionSet"
 215         else:
 216             sorted_sets = []
 217             for selset in arm.selection_sets:
 218                 if selset.name.startswith("SelectionSet."):
 219                     index = selset.name[13:]
 220                     if index.isdigit():
 221                         sorted_sets.append(index)
 222             sorted_sets = sorted(sorted_sets)
 223             min_index = 1
 224             for num in sorted_sets:
 225                 num = int(num)
 226                 if min_index < num:
 227                     break
 228                 min_index = num + 1
 229             new_sel_set.name = "SelectionSet.{:03d}".format(min_index)
 230
 231         # select newly created set
 232         arm.active_selection_set = len(arm.selection_sets) - 1
 233
 234         return {'FINISHED'}
 235
 236
 237 class POSE_OT_selection_set_remove(NeedSelSetPluginOperator):
 238     bl_idname = "pose.selection_set_remove"
 239     bl_label = "Delete Selection Set"
 240     bl_description = "Delete a Selection Set"
 241     bl_options = {'UNDO', 'REGISTER'}
 242
 243     def execute(self, context):
 244         arm = context.object
 245
 246         arm.selection_sets.remove(arm.active_selection_set)
 247
 248         # change currently active selection set
 249         numsets = len(arm.selection_sets)
 250         if (arm.active_selection_set > (numsets - 1) and numsets > 0):
 251             arm.active_selection_set = len(arm.selection_sets) - 1
 252
 253         return {'FINISHED'}
 254
 255
 256 class POSE_OT_selection_set_assign(PluginOperator):
 257     bl_idname = "pose.selection_set_assign"
 258     bl_label = "Add Bones to Selection Set"
 259     bl_description = "Add selected bones to Selection Set"
 260     bl_options = {'UNDO', 'REGISTER'}
 261
 262     def invoke(self, context, event):
 263         arm = context.object
 264
 265         if not (arm.active_selection_set < len(arm.selection_sets)):
 266             bpy.ops.wm.call_menu("INVOKE_DEFAULT",
 267                 name="pose.selection_set_create_new_popup")
 268         else:
 269             bpy.ops.pose.selection_set_assign('EXEC_DEFAULT')
 270
 271         return {'FINISHED'}
 272
 273
 274     def execute(self, context):
 275         arm = context.object
 276         act_sel_set = arm.selection_sets[arm.active_selection_set]
 277
 278         # iterate only the selected bones in current pose that are not hidden
 279         for bone in context.selected_pose_bones:
 280             if bone.name not in act_sel_set.bone_ids:
 281                 bone_id = act_sel_set.bone_ids.add()
 282                 bone_id.name = bone.name
 283
 284         return {'FINISHED'}
 285
 286
 287 class POSE_OT_selection_set_unassign(NeedSelSetPluginOperator):
 288     bl_idname = "pose.selection_set_unassign"
 289     bl_label = "Remove Bones from Selection Set"
 290     bl_description = "Remove selected bones from Selection Set"
 291     bl_options = {'UNDO', 'REGISTER'}
 292
 293     def execute(self, context):
 294         arm = context.object
 295         act_sel_set = arm.selection_sets[arm.active_selection_set]
 296
 297         # iterate only the selected bones in current pose that are not hidden
 298         for bone in context.selected_pose_bones:
 299             if bone.name in act_sel_set.bone_ids:
 300                 idx = act_sel_set.bone_ids.find(bone.name)
 301                 act_sel_set.bone_ids.remove(idx)
 302
 303         return {'FINISHED'}
 304
 305
 306 class POSE_OT_selection_set_select(NeedSelSetPluginOperator):
 307     bl_idname = "pose.selection_set_select"
 308     bl_label = "Select Selection Set"
 309     bl_description = "Add Selection Set bones to current selection"
 310     bl_options = {'UNDO', 'REGISTER'}
 311
 312     def execute(self, context):
 313         arm = context.object
 314         act_sel_set = arm.selection_sets[arm.active_selection_set]
 315
 316         for bone in context.visible_pose_bones:
 317             if bone.name in act_sel_set.bone_ids:
 318                  bone.bone.select = True
 319
 320         return {'FINISHED'}
 321
 322
 323 class POSE_OT_selection_set_deselect(NeedSelSetPluginOperator):
 324     bl_idname = "pose.selection_set_deselect"
 325     bl_label = "Deselect Selection Set"
 326     bl_description = "Remove Selection Set bones from current selection"
 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         for bone in context.selected_pose_bones:
 334             if bone.name in act_sel_set.bone_ids:
 335                 bone.bone.select = False
 336
 337         return {'FINISHED'}
 338
 339
 340 class POSE_OT_selection_set_add_and_assign(PluginOperator):
 341     bl_idname = "pose.selection_set_add_and_assign"
 342     bl_label = "Create and Add Bones to Selection Set"
 343     bl_description = "Creates a new Selection Set with the currently selected bones"
 344     bl_options = {'UNDO', 'REGISTER'}
 345
 346     def execute(self, context):
 347         bpy.ops.pose.selection_set_add('EXEC_DEFAULT')
 348         bpy.ops.pose.selection_set_assign('EXEC_DEFAULT')
 349         return {'FINISHED'}
 350
 351 # Registry ####################################################################
 352
 353 classes = (
 354     POSE_MT_create_new_selection_set,
 355     POSE_MT_selection_sets_specials,
 356     POSE_PT_selection_sets,
 357     POSE_UL_selection_set,
 358     SelectionEntry,
 359     SelectionSet,
 360     POSE_OT_selection_set_move,
 361     POSE_OT_selection_set_add,
 362     POSE_OT_selection_set_remove,
 363     POSE_OT_selection_set_assign,
 364     POSE_OT_selection_set_unassign,
 365     POSE_OT_selection_set_select,
 366     POSE_OT_selection_set_deselect,
 367     POSE_OT_selection_set_add_and_assign,
 368 )
 369
 370 def register():
 371     for cls in classes:
 372         bpy.utils.register_class(cls)
 373
 374     bpy.types.Object.selection_sets = CollectionProperty(
 375         type=SelectionSet,
 376         name="Selection Sets",
 377         description="List of groups of bones for easy selection"
 378     )
 379     bpy.types.Object.active_selection_set = IntProperty(
 380         name="Active Selection Set",
 381         description="Index of the currently active selection set",
 382         default=0
 383     )
 384
 385
 386 def unregister():
 387     for cls in classes:
 388         bpy.utils.unregister_class(cls)
 389
 390     del bpy.types.Object.selection_sets
 391     del bpy.types.Object.active_selection_set
 392
 393
 394 if __name__ == "__main__":
 395     register()