anim_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, see <http://www.gnu.org/licenses/>
  15  #  and write to the Free Software Foundation, Inc., 51 Franklin Street,
  16  #  Fifth Floor, Boston, MA  02110-1301, USA..
  17  #
  18  #  The Original Code is Copyright (C) 2012 Blender Foundation ###
  19  #  All rights reserved.
  20  #
  21  #
  22  #  The Original Code is: all of this file.
  23  #
  24  #  Contributor(s): Dan Eicher.
  25  #
  26  #  ***** END GPL LICENSE BLOCK *****
  27
  28 # <pep8 compliant>
  29
  30 import string
  31 import bpy
  32
  33 bl_info = {
  34   "name": "Selection Set",
  35   "author": "Dan Eicher",
  36   "version": (0, 1, 0),
  37   "blender": (2, 6, 3),
  38   "location": "Properties -> Object Data -> Selection Sets",
  39   "description": "Selection Sets to select groups of bones",
  40   "warning": "Proxy armatures need to export sets and run generated script on re-opening file",
  41   "wiki_url": "http://wiki.blender.org/index.php/Extensions:2.6/Py/Scripts/Animation/SelectionSets",
  42   "tracker_url": "http://projects.blender.org/tracker/index.php?func=detail&aid=31492",
  43   "category": "Animation"
  44 }
  45
  46
  47 script_template = '''# generated by ANIM_OT_selection_set_export -- abandon all hope, ye who hand edit!
  48
  49 import bpy
  50
  51 def selection_set_import():
  52     arm = bpy.data.armatures['${name}']
  53
  54     set_list = [${set_list}
  55                ]
  56
  57     for name, bones in set_list:
  58         if arm.selection_sets.find(name) == -1:
  59             sel_set = arm.selection_sets.add()
  60             sel_set.name = name
  61
  62             for bone in bones:
  63                 set_bone = sel_set.bones.add()
  64                 set_bone.name = bone
  65
  66
  67 if __name__ == "__main__":
  68     selection_set_import()
  69 '''
  70
  71
  72 def generic_poll(context):
  73     return (context.mode == 'POSE' and context.object and context.object.type == 'ARMATURE')
  74
  75
  76 def active_selection_set_update_func(self, context):
  77     idx = self.active_selection_set
  78     if idx < -1 or idx >= len(self.selection_sets):
  79         self.active_selection_set = -1
  80         raise IndexError('Armature.active_selection_set: out of range')
  81
  82
  83 class SelectionSetBone(bpy.types.PropertyGroup):
  84     name = bpy.props.StringProperty()
  85
  86
  87 class SelectionSet(bpy.types.PropertyGroup):
  88     name = bpy.props.StringProperty(options={'LIBRARY_EDITABLE'})
  89     bones = bpy.props.CollectionProperty(type=SelectionSetBone)
  90
  91
  92 class ANIM_OT_selection_set_add(bpy.types.Operator):
  93     """Add a new selection set"""
  94     bl_idname = "anim.selection_set_add"
  95     bl_label = "Selection Set Add"
  96
  97     @classmethod
  98     def poll(cls, context):
  99         return generic_poll(context)
 100
 101     def execute(self, context):
 102         arm = context.active_object.data
 103
 104         tmp_name = name = 'Set'
 105         name_sub = 1
 106         while arm.selection_sets.find(name) != -1:
 107             name = tmp_name + ' ' + str(name_sub)
 108             name_sub += 1
 109
 110         sel_set = arm.selection_sets.add()
 111         sel_set.name = name
 112
 113         arm.active_selection_set = arm.selection_sets.find(name)
 114
 115         return {'FINISHED'}
 116
 117
 118 class ANIM_OT_selection_set_remove(bpy.types.Operator):
 119     """Remove the active selection set"""
 120     bl_idname = "anim.selection_set_remove"
 121     bl_label = "Selection Set Remove"
 122
 123     @classmethod
 124     def poll(cls, context):
 125         arm = context.active_object.data
 126         return (generic_poll(context) and arm.active_selection_set != -1)
 127
 128     def execute(self, context):
 129         arm = context.active_object.data
 130         active_index = arm.active_selection_set
 131
 132         arm.selection_sets.remove(active_index)
 133
 134         if active_index >= len(arm.selection_sets):
 135             arm.active_selection_set = len(arm.selection_sets) - 1
 136
 137         return {'FINISHED'}
 138
 139
 140 class ANIM_OT_selection_set_assign(bpy.types.Operator):
 141     """Add selected bones to the active selection set"""
 142     bl_idname = "anim.selection_set_assign"
 143     bl_label = "Selection Set Assign"
 144
 145     @classmethod
 146     def poll(cls, context):
 147         arm = context.active_object.data
 148         return (generic_poll(context) and arm.active_selection_set != -1)
 149
 150     def execute(self, context):
 151         arm = context.active_object.data
 152         sel_set = arm.selection_sets[arm.active_selection_set]
 153         bones = [bone for bone in arm.bones if bone.select]
 154
 155         for bone in bones:
 156             if sel_set.bones.find(bone.name) == -1:
 157                 set_bone = sel_set.bones.add()
 158                 set_bone.name = bone.name
 159
 160         return {'FINISHED'}
 161
 162
 163 class ANIM_OT_selection_set_unassign(bpy.types.Operator):
 164     """Remove selected bones from the active selection set"""
 165     bl_idname = "anim.selection_set_unassign"
 166     bl_label = "Selection Set Unassign"
 167
 168     @classmethod
 169     def poll(cls, context):
 170         arm = context.active_object.data
 171         return (generic_poll(context) and arm.active_selection_set != -1)
 172
 173     def execute(self, context):
 174         arm = context.active_object.data
 175         sel_set = arm.selection_sets[arm.active_selection_set]
 176         bones = [bone for bone in arm.bones if bone.select]
 177
 178         for bone in bones:
 179             bone_index = sel_set.bones.find(bone.name)
 180             if bone_index != -1:
 181                 sel_set.bones.remove(bone_index)
 182
 183         return {'FINISHED'}
 184
 185
 186 class ANIM_OT_selection_set_select(bpy.types.Operator):
 187     """Select bones in selection set"""
 188     bl_idname = "anim.selection_set_select"
 189     bl_label = "Selection Set Select Bones"
 190
 191     @classmethod
 192     def poll(cls, context):
 193         arm = context.active_object.data
 194         return (generic_poll(context) and arm.active_selection_set != -1)
 195
 196     def execute(self, context):
 197         arm = context.active_object.data
 198         sel_set = arm.selection_sets[arm.active_selection_set]
 199
 200         for bone in sel_set.bones:
 201             try:
 202                 arm.bones[bone.name].select = True
 203             except:
 204                 bone_index = sel_set.bones.find(bone.name)
 205                 sel_set.bones.remove(bone_index)
 206
 207         return {'FINISHED'}
 208
 209
 210 class ANIM_OT_selection_set_deselect(bpy.types.Operator):
 211     """Deselect bones in selection set"""
 212     bl_idname = "anim.selection_set_deselect"
 213     bl_label = "Selection Set Deselect Bones"
 214
 215     @classmethod
 216     def poll(cls, context):
 217         arm = context.active_object.data
 218         return (generic_poll(context) and arm.active_selection_set != -1)
 219
 220     def execute(self, context):
 221         arm = context.active_object.data
 222         sel_set = arm.selection_sets[arm.active_selection_set]
 223
 224         for bone in sel_set.bones:
 225             try:
 226                 arm.bones[bone.name].select = False
 227             except:
 228                 bone_index = sel_set.bones.find(bone.name)
 229                 sel_set.bones.remove(bone_index)
 230
 231         return {'FINISHED'}
 232
 233
 234 class ANIM_OT_selection_set_export(bpy.types.Operator):
 235     """Export selection set data to a python script"""
 236     bl_idname = "anim.selection_set_export"
 237     bl_label = "Selection Set Export"
 238
 239     @classmethod
 240     def poll(cls, context):
 241         return generic_poll(context)
 242
 243     def execute(self, context):
 244         arm = context.active_object.data
 245         set_script = string.Template(script_template)
 246         set_list = ""
 247
 248         for sel_set in arm.selection_sets:
 249             set_bones = ""
 250             for bone in sel_set.bones:
 251                 set_bones += "'" + bone.name + "',"
 252             set_list += "\n                ('{name}', [{bones}]),".format(name=sel_set.name, bones=set_bones)
 253
 254         try:
 255             script_file = bpy.data.texts['{arm.name}SelectionSetImport.py'.format(arm=arm)]
 256         except:
 257             script_file = bpy.data.texts.new('{arm.name}SelectionSetImport.py'.format(arm=arm))
 258
 259         script_file.clear()
 260         script_file.write(set_script.substitute(name=arm.name, set_list=set_list))
 261
 262         return {'FINISHED'}
 263
 264
 265 class DATA_PT_bone_sets(bpy.types.Panel):
 266     bl_space_type = 'PROPERTIES'
 267     bl_region_type = 'WINDOW'
 268     bl_context = "data"
 269     bl_label = "Selection Sets"
 270
 271     @classmethod
 272     def poll(cls, context):
 273         return (context.object and context.object.type == 'ARMATURE' and context.object.pose)
 274
 275     def draw(self, context):
 276         layout = self.layout
 277         ob = context.object
 278         arm = ob.data
 279         sel_set = None
 280
 281         if arm.active_selection_set != -1:
 282             try:
 283                 sel_set = arm.selection_sets[arm.active_selection_set]
 284             except:
 285                 pass
 286
 287         row = layout.row()
 288
 289         row.template_list(arm, "selection_sets", arm, "active_selection_set",
 290                           rows=(5 if len(arm.selection_sets) else 2))
 291
 292         col = row.column(align=True)
 293         col.operator("anim.selection_set_add", icon='ZOOMIN', text="")
 294         col.operator("anim.selection_set_remove", icon='ZOOMOUT', text="")
 295
 296         if sel_set:
 297             col = layout.column()
 298             col.prop(sel_set, "name", text="Name")
 299
 300         row = layout.row()
 301
 302         sub = row.row(align=True)
 303         sub.operator("anim.selection_set_assign", text="Assign")
 304         sub.operator("anim.selection_set_unassign", text="Remove")
 305
 306         sub = row.row(align=True)
 307         sub.operator("anim.selection_set_select", text="Select")
 308         sub.operator("anim.selection_set_deselect", text="Deselect")
 309
 310         row = layout.row()
 311         row.operator("anim.selection_set_export", text="Export Selection Sets")
 312
 313
 314 def register():
 315     bpy.utils.register_class(SelectionSetBone)
 316     bpy.utils.register_class(SelectionSet)
 317     bpy.utils.register_class(ANIM_OT_selection_set_add)
 318     bpy.utils.register_class(ANIM_OT_selection_set_remove)
 319     bpy.utils.register_class(ANIM_OT_selection_set_assign)
 320     bpy.utils.register_class(ANIM_OT_selection_set_unassign)
 321     bpy.utils.register_class(ANIM_OT_selection_set_select)
 322     bpy.utils.register_class(ANIM_OT_selection_set_deselect)
 323     bpy.utils.register_class(ANIM_OT_selection_set_export)
 324     bpy.utils.register_class(DATA_PT_bone_sets)
 325
 326     bpy.types.Armature.selection_sets = bpy.props.CollectionProperty(type=SelectionSet,
 327                                             name="Selection Sets",
 328                                             description="Collection of bones to be selected")
 329
 330     bpy.types.Armature.active_selection_set = bpy.props.IntProperty(name="Active Selection Set",
 331                                                                     default=-1,
 332                                                                     options={'LIBRARY_EDITABLE'},
 333                                                                     update=active_selection_set_update_func)
 334
 335
 336 def unregister():
 337     del bpy.types.Armature.selection_sets
 338     del bpy.types.Armature.active_selection_set
 339
 340     bpy.utils.unregister_class(SelectionSet)
 341     bpy.utils.unregister_class(SelectionSetBone)
 342     bpy.utils.unregister_class(ANIM_OT_selection_set_add)
 343     bpy.utils.unregister_class(ANIM_OT_selection_set_remove)
 344     bpy.utils.unregister_class(ANIM_OT_selection_set_assign)
 345     bpy.utils.unregister_class(ANIM_OT_selection_set_unassign)
 346     bpy.utils.unregister_class(ANIM_OT_selection_set_select)
 347     bpy.utils.unregister_class(ANIM_OT_selection_set_deselect)
 348     bpy.utils.unregister_class(ANIM_OT_selection_set_export)
 349     bpy.utils.unregister_class(DATA_PT_bone_sets)
 350
 351 if __name__ == "__main__":
 352     register()
 353