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