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