27239b0e31da352bbaba8bc3cfae40442a895e97
[blender-addons-contrib.git] / 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