0dbd964acb596c08c2bf27a6e338b0ea11ae5288
[blender-staging.git] / release / scripts / ui / properties_data_armature_rigify.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 LICENSE BLOCK #####
18
19 # <pep8 compliant>
20 import bpy
21
22
23 class PoseTemplateSettings(bpy.types.IDPropertyGroup):
24     pass
25
26
27 class PoseTemplate(bpy.types.IDPropertyGroup):
28     pass
29
30
31 def metarig_templates():
32     import rigify
33     return rigify.get_submodule_types()
34
35
36 class DATA_PT_template(bpy.types.Panel):
37     bl_label = "Meta-Rig Templates"
38     bl_space_type = 'PROPERTIES'
39     bl_region_type = 'WINDOW'
40     bl_context = "data"
41     bl_default_closed = True
42
43     templates = []
44
45     @classmethod
46     def poll(cls, context):
47         if not context.armature:
48             return False
49         obj = context.object
50         if obj:
51             return (obj.mode in ('POSE', 'OBJECT'))
52         return False
53
54     def draw(self, context):
55         layout = self.layout
56         try:
57             active_type = context.active_pose_bone["type"]
58         except:
59             active_type = None
60
61         scene = context.scene
62         pose_templates = scene.pose_templates
63
64         if pose_templates.active_template_index == -1:
65             pose_templates.active_template_index = 0
66
67         if not self.templates:
68             self.templates[:] = metarig_templates()
69
70         while(len(pose_templates.templates) < len(self.templates)):
71             pose_templates.templates.add()
72         while(len(pose_templates.templates) > len(self.templates)):
73             pose_templates.templates.remove(0)
74
75         for i, template_name in enumerate(self.templates):
76             template = pose_templates.templates[i]
77             if active_type == template_name:
78                 template.name = "<%s>" % template_name.replace("_", " ")
79             else:
80                 template.name = " %s " % template_name.replace("_", " ")
81
82         row = layout.row()
83         row.operator("pose.metarig_generate", text="Generate")
84         row.operator("pose.metarig_validate", text="Check")
85         row.operator("pose.metarig_graph", text="Graph")
86         row = layout.row()
87         row.prop(pose_templates, "generate_def_rig")
88
89         row = layout.row()
90         col = row.column()
91         col.template_list(pose_templates, "templates", pose_templates, "active_template_index", rows=1)
92
93         subrow = col.split(percentage=0.5, align=True)
94         subsubrow = subrow.row(align=True)
95         subsubrow.operator("pose.metarig_assign", text="Assign")
96         subsubrow.operator("pose.metarig_clear", text="Clear")
97
98         if self.templates:
99             subsubrow = subrow.split(percentage=0.8)
100             subsubrow.operator("pose.metarig_sample_add", text="Sample").metarig_type = self.templates[pose_templates.active_template_index]
101             subsubrow.operator("pose.metarig_sample_add", text="All").metarig_type = "" # self.templates[pose_templates.active_template_index]
102
103             sub = row.column(align=True)
104             sub.operator("pose.metarig_reload", icon="FILE_REFRESH", text="")
105
106
107 # operators
108 from bpy.props import StringProperty
109
110
111 class Reload(bpy.types.Operator):
112     '''Re-Scan the metarig package directory for scripts'''
113
114     bl_idname = "pose.metarig_reload"
115     bl_label = "Re-Scan the list of metarig types"
116
117     def execute(self, context):
118         DATA_PT_template.templates[:] = metarig_templates()
119         return {'FINISHED'}
120
121
122 def rigify_report_exception(operator, exception):
123     import traceback
124     import sys
125     import os
126     # find the module name where the error happened
127     # hint, this is the metarig type!
128     exceptionType, exceptionValue, exceptionTraceback = sys.exc_info()
129     fn = traceback.extract_tb(exceptionTraceback)[-1][0]
130     fn = os.path.basename(fn)
131     fn = os.path.splitext(fn)[0]
132     message = []
133     if fn.startswith("__"):
134         message.append("Incorrect armature...")
135     else:
136         message.append("Incorrect armature for type '%s'" % fn)
137     message.append(exception.message)
138
139     message.reverse() # XXX - stupid! menu's are upside down!
140
141     operator.report(set(['INFO']), '\n'.join(message))
142
143
144 class Generate(bpy.types.Operator):
145     '''Generates a metarig from the active armature'''
146
147     bl_idname = "pose.metarig_generate"
148     bl_label = "Generate Metarig"
149
150     def execute(self, context):
151         import rigify
152         reload(rigify)
153
154         meta_def = context.scene.pose_templates.generate_def_rig
155
156         try:
157             rigify.generate_rig(context, context.object, META_DEF=meta_def)
158         except rigify.RigifyError as rig_exception:
159             rigify_report_exception(self, rig_exception)
160
161         return {'FINISHED'}
162
163
164 class Validate(bpy.types.Operator):
165     '''Validate a metarig from the active armature'''
166
167     bl_idname = "pose.metarig_validate"
168     bl_label = "Validate Metarig"
169
170     def execute(self, context):
171         import rigify
172         reload(rigify)
173         try:
174             rigify.validate_rig(context, context.object)
175         except rigify.RigifyError as rig_exception:
176             rigify_report_exception(self, rig_exception)
177         return {'FINISHED'}
178
179
180 class Sample(bpy.types.Operator):
181     '''Create a sample metarig to be modified before generating the final rig.'''
182
183     bl_idname = "pose.metarig_sample_add"
184     bl_label = "Re-Scan Metarig Scripts"
185
186     metarig_type = StringProperty(name="Type", description="Name of the rig type to generate a sample of, a blank string for all", maxlen=128, default="")
187
188     def execute(self, context):
189         import rigify
190         reload(rigify)
191         final = (self.properties.metarig_type == "")
192         objects = rigify.generate_test(context, metarig_type=self.properties.metarig_type, GENERATE_FINAL=final)
193
194         if len(objects) > 1:
195             for i, (obj_meta, obj_gen) in enumerate(objects):
196                 obj_meta.location.x = i * 1.0
197                 if obj_gen:
198                     obj_gen.location.x = i * 1.0
199
200         return {'FINISHED'}
201
202
203 class Graph(bpy.types.Operator):
204     '''Create a graph from the active armature through graphviz'''
205
206     bl_idname = "pose.metarig_graph"
207     bl_label = "Pose Graph"
208
209     def execute(self, context):
210         import os
211         import graphviz_export
212         import bpy
213         reload(graphviz_export)
214         obj = bpy.context.object
215         path = os.path.splitext(bpy.data.filepath)[0] + "-" + bpy.path.clean_name(obj.name)
216         path_dot = path + ".dot"
217         path_png = path + ".png"
218         saved = graphviz_export.graph_armature(bpy.context.object, path_dot, CONSTRAINTS=False, DRIVERS=False)
219
220         if saved:
221             # if we seriously want this working everywhere we'll need some new approach
222             os.system("dot -Tpng %s > %s; gnome-open %s &" % (path_dot, path_png, path_png))
223             #os.system("python /b/xdot.py '%s' &" % path_dot)
224
225         return {'FINISHED'}
226
227
228 class AsScript(bpy.types.Operator):
229     '''Write the edit armature out as a python script'''
230
231     bl_idname = "pose.metarig_to_script"
232     bl_label = "Write Metarig to Script"
233     bl_options = {'REGISTER', 'UNDO'}
234
235     filepath = StringProperty(name="File Path", description="File path used for exporting the Armature file", maxlen=1024, default="")
236
237     def execute(self, context):
238         import rigify_utils
239         reload(rigify_utils)
240         obj = context.object
241         code = rigify_utils.write_meta_rig(obj)
242         path = self.properties.filepath
243         file = open(path, "w")
244         file.write(code)
245         file.close()
246
247         return {'FINISHED'}
248
249     def invoke(self, context, event):
250         import os
251         obj = context.object
252         self.properties.filepath = os.path.splitext(bpy.data.filepath)[0] + "-" + bpy.path.clean_name(obj.name) + ".py"
253         wm = context.manager
254         wm.add_fileselect(self)
255         return {'RUNNING_MODAL'}
256
257
258 # operators that use the GUI
259 class ActiveAssign(bpy.types.Operator):
260     '''Assign to the active posebone'''
261
262     bl_idname = "pose.metarig_assign"
263     bl_label = "Assign to the active posebone"
264
265     @classmethod
266     def poll(cls, context):
267         bone = context.active_pose_bone
268         return bool(bone and bone.id_data.mode == 'POSE')
269
270     def execute(self, context):
271         scene = context.scene
272         pose_templates = scene.pose_templates
273         template_name = DATA_PT_template.templates[pose_templates.active_template_index]
274         context.active_pose_bone["type"] = template_name
275         return {'FINISHED'}
276
277
278 class ActiveClear(bpy.types.Operator):
279     '''Clear type from the active posebone'''
280
281     bl_idname = "pose.metarig_clear"
282     bl_label = "Metarig Clear Type"
283
284     @classmethod
285     def poll(cls, context):
286         bone = context.active_pose_bone
287         return bool(bone and bone.id_data.mode == 'POSE')
288
289     def execute(self, context):
290         scene = context.scene
291         try:
292             del context.active_pose_bone["type"]
293         except:
294             return {'CANCELLED'}
295         return {'FINISHED'}
296
297
298 class INFO_MT_armature_metarig_add(bpy.types.Menu):
299     bl_idname = "INFO_MT_armature_metarig_add"
300     bl_label = "Meta-Rig"
301
302     def draw(self, context):
303         import rigify
304
305         layout = self.layout
306         layout.operator_context = 'INVOKE_REGION_WIN'
307
308         for submodule_type in rigify.get_submodule_types():
309             text = bpy.path.display_name(submodule_type)
310             layout.operator("pose.metarig_sample_add", text=text, icon='OUTLINER_OB_ARMATURE').metarig_type = submodule_type
311
312 menu_func = (lambda self, context: self.layout.menu("INFO_MT_armature_metarig_add", icon='OUTLINER_OB_ARMATURE'))
313 import space_info # ensure the menu is loaded first
314
315
316 def register():
317     PoseTemplate.StringProperty(attr="name",
318                     name="Name of the slave",
319                     description="",
320                     maxlen=64,
321                     default="")
322
323
324     PoseTemplateSettings.CollectionProperty(attr="templates", type=PoseTemplate, name="Templates", description="")
325     PoseTemplateSettings.IntProperty(attr="active_template_index",
326                     name="Index of the active slave",
327                     description="",
328                     default=-1,
329                     min=-1,
330                     max=65535)
331
332     PoseTemplateSettings.BoolProperty(attr="generate_def_rig",
333                     name="Create Deform Rig",
334                     description="Create a copy of the metarig, constrainted by the generated rig",
335                     default=False)
336
337     bpy.types.Scene.PointerProperty(attr="pose_templates", type=PoseTemplateSettings, name="Pose Templates", description="Pose Template Settings")
338
339     space_info.INFO_MT_armature_add.append(menu_func)
340
341
342 def unregister():
343     bpy.types.Scene.RemoveProperty("pose_templates")
344     space_info.INFO_MT_armature_add.remove(menu_func)
345
346 if __name__ == "__main__":
347     register()