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
21import bpy
22from bpy.props import StringProperty
23
24from .utils import get_rig_type, MetarigError
25from .utils import write_metarig, write_widget
26from . import rig_lists
27from . import generate
28
29
30class DATA_PT_rigify_buttons(bpy.types.Panel):
31    bl_label = "Rigify Buttons"
32    bl_space_type = 'PROPERTIES'
33    bl_region_type = 'WINDOW'
34    bl_context = "data"
35    #bl_options = {'DEFAULT_OPEN'}
36
37    @classmethod
38    def poll(cls, context):
39        if not context.armature:
40            return False
41        #obj = context.object
42        #if obj:
43        #    return (obj.mode in {'POSE', 'OBJECT', 'EDIT'})
44        #return False
45        return True
46
47    def draw(self, context):
48        C = context
49        layout = self.layout
50        obj = context.object
51        id_store = C.window_manager
52
53        if obj.mode in {'POSE', 'OBJECT'}:
54            layout.operator("pose.rigify_generate", text="Generate")
55        elif obj.mode == 'EDIT':
56            # Build types list
57            collection_name = str(id_store.rigify_collection).replace(" ", "")
58
59            for i in range(0, len(id_store.rigify_types)):
60                id_store.rigify_types.remove(0)
61
62            for r in rig_lists.rig_list:
63                # collection = r.split('.')[0]  # UNUSED
64                if collection_name == "All":
65                    a = id_store.rigify_types.add()
66                    a.name = r
67                elif r.startswith(collection_name + '.'):
68                    a = id_store.rigify_types.add()
69                    a.name = r
70                elif (collection_name == "None") and ("." not in r):
71                    a = id_store.rigify_types.add()
72                    a.name = r
73
74            ## Rig collection field
75            #row = layout.row()
76            #row.prop(id_store, 'rigify_collection', text="Category")
77
78            # Rig type list
79            row = layout.row()
80            row.template_list("UI_UL_list", "rigify_types", id_store, "rigify_types", id_store, 'rigify_active_type')
81
82            props = layout.operator("armature.metarig_sample_add", text="Add sample")
83            props.metarig_type = id_store.rigify_types[id_store.rigify_active_type].name
84
85
86class DATA_PT_rigify_layer_names(bpy.types.Panel):
87    bl_label = "Rigify Layer Names"
88    bl_space_type = 'PROPERTIES'
89    bl_region_type = 'WINDOW'
90    bl_context = "data"
91    bl_options = {'DEFAULT_CLOSED'}
92
93    @classmethod
94    def poll(cls, context):
95        if not context.armature:
96            return False
97        return True
98
99    def draw(self, context):
100        layout = self.layout
101        obj = context.object
102        arm = obj.data
103
104        # Ensure that the layers exist
105        if 0:
106            for i in range(1 + len(arm.rigify_layers), 29):
107                arm.rigify_layers.add()
108        else:
109            # Can't add while drawing, just use button
110            if len(arm.rigify_layers) < 28:
111                layout.operator("pose.rigify_layer_init")
112                return
113
114        # UI
115        for i, rigify_layer in enumerate(arm.rigify_layers):
116            # note: rigify_layer == arm.rigify_layers[i]
117            if (i % 16) == 0:
118                col = layout.column()
119                if i == 0:
120                    col.label(text="Top Row:")
121                else:
122                    col.label(text="Bottom Row:")
123            if (i % 8) == 0:
124                col = layout.column(align=True)
125            row = col.row()
126            row.prop(arm, "layers", index=i, text="", toggle=True)
127            split = row.split(factor=0.8)
128            split.prop(rigify_layer, "name",  text="Layer %d" % (i + 1))
129            split.prop(rigify_layer, "row",   text="")
130
131            #split.prop(rigify_layer, "column", text="")
132
133
134class BONE_PT_rigify_buttons(bpy.types.Panel):
135    bl_label = "Rigify Type"
136    bl_space_type = 'PROPERTIES'
137    bl_region_type = 'WINDOW'
138    bl_context = "bone"
139    #bl_options = {'DEFAULT_OPEN'}
140
141    @classmethod
142    def poll(cls, context):
143        if not context.armature or not context.active_pose_bone:
144            return False
145        obj = context.object
146        if obj:
147            return obj.mode == 'POSE'
148        return False
149
150    def draw(self, context):
151        C = context
152        id_store = C.window_manager
153        bone = context.active_pose_bone
154        collection_name = str(id_store.rigify_collection).replace(" ", "")
155        rig_name = str(context.active_pose_bone.rigify_type).replace(" ", "")
156
157        layout = self.layout
158
159        # Build types list
160        for i in range(0, len(id_store.rigify_types)):
161            id_store.rigify_types.remove(0)
162
163        for r in rig_lists.rig_list:
164            # collection = r.split('.')[0]  # UNUSED
165            if collection_name == "All":
166                a = id_store.rigify_types.add()
167                a.name = r
168            elif r.startswith(collection_name + '.'):
169                a = id_store.rigify_types.add()
170                a.name = r
171            elif collection_name == "None" and len(r.split('.')) == 1:
172                a = id_store.rigify_types.add()
173                a.name = r
174
175        # Rig type field
176        row = layout.row()
177        row.prop_search(bone, "rigify_type", id_store, "rigify_types", text="Rig type:")
178
179        # Rig type parameters / Rig type non-exist alert
180        if rig_name != "":
181            try:
182                rig = get_rig_type(rig_name)
183                rig.Rig
184            except (ImportError, AttributeError):
185                row = layout.row()
186                box = row.box()
187                box.label(text="ALERT: type \"%s\" does not exist!" % rig_name)
188            else:
189                try:
190                    rig.parameters_ui
191                except AttributeError:
192                    col = layout.column()
193                    col.label(text="No options")
194                else:
195                    col = layout.column()
196                    col.label(text="Options:")
197                    box = layout.box()
198                    rig.parameters_ui(box, bone.rigify_parameters)
199
200
201class VIEW3D_PT_tools_rigify_dev(bpy.types.Panel):
202    bl_label = "Rigify Dev Tools"
203    bl_space_type = 'VIEW_3D'
204    bl_region_type = 'UI'
205    bl_category = 'Rigify'
206
207    @classmethod
208    def poll(cls, context):
209        return context.active_object is not None and context.mode in {'EDIT_ARMATURE','EDIT_MESH'}
210
211    def draw(self, context):
212        obj = context.active_object
213        if obj is not None:
214            if context.mode == 'EDIT_ARMATURE':
215                r = self.layout.row()
216                r.operator("armature.rigify_encode_metarig", text="Encode Metarig to Python")
217                r = self.layout.row()
218                r.operator("armature.rigify_encode_metarig_sample", text="Encode Sample to Python")
219
220            if context.mode == 'EDIT_MESH':
221                r = self.layout.row()
222                r.operator("mesh.rigify_encode_mesh_widget", text="Encode Mesh Widget to Python")
223
224#~ class VIEW3D_MT_armature_metarig_add(bpy.types.Menu):
225    #~ bl_idname = "VIEW3D_MT_armature_metarig_add"
226    #~ bl_label = "Meta-Rig"
227
228    #~ def draw(self, context):
229        #~ import rigify
230
231        #~ layout = self.layout
232        #~ layout.operator_context = 'INVOKE_REGION_WIN'
233
234        #~ for submodule_type in rigify.get_submodule_types():
235            #~ text = bpy.path.display_name(submodule_type)
236            #~ layout.operator("pose.metarig_sample_add", text=text, icon='OUTLINER_OB_ARMATURE').metarig_type = submodule_type
237
238
239def rigify_report_exception(operator, exception):
240    import traceback
241    import sys
242    import os
243    # find the module name where the error happened
244    # hint, this is the metarig type!
245    exceptionType, exceptionValue, exceptionTraceback = sys.exc_info()
246    fn = traceback.extract_tb(exceptionTraceback)[-1][0]
247    fn = os.path.basename(fn)
248    fn = os.path.splitext(fn)[0]
249    message = []
250    if fn.startswith("__"):
251        message.append("Incorrect armature...")
252    else:
253        message.append("Incorrect armature for type '%s'" % fn)
254    message.append(exception.message)
255
256    message.reverse()  # XXX - stupid! menu's are upside down!
257
258    operator.report({'INFO'}, '\n'.join(message))
259
260
261class LayerInit(bpy.types.Operator):
262    """Initialize armature rigify layers"""
263
264    bl_idname = "pose.rigify_layer_init"
265    bl_label = "Add Rigify Layers"
266    bl_options = {'UNDO'}
267
268    def execute(self, context):
269        obj = context.object
270        arm = obj.data
271        for i in range(1 + len(arm.rigify_layers), 29):
272            arm.rigify_layers.add()
273        return {'FINISHED'}
274
275
276class Generate(bpy.types.Operator):
277    """Generates a rig from the active metarig armature"""
278
279    bl_idname = "pose.rigify_generate"
280    bl_label = "Rigify Generate Rig"
281    bl_options = {'UNDO'}
282
283    def execute(self, context):
284        import importlib
285        importlib.reload(generate)
286
287        try:
288            generate.generate_rig(context, context.object)
289        except MetarigError as rig_exception:
290            rigify_report_exception(self, rig_exception)
291
292        return {'FINISHED'}
293
294
295class Sample(bpy.types.Operator):
296    """Create a sample metarig to be modified before generating """ \
297    """the final rig"""
298
299    bl_idname = "armature.metarig_sample_add"
300    bl_label = "Add a sample metarig for a rig type"
301    bl_options = {'UNDO'}
302
303    metarig_type: StringProperty(
304        name="Type",
305        description="Name of the rig type to generate a sample of",
306        maxlen=128,
307    )
308
309    def execute(self, context):
310        if context.mode == 'EDIT_ARMATURE' and self.metarig_type != "":
311            try:
312                rig = get_rig_type(self.metarig_type)
313                create_sample = rig.create_sample
314            except (ImportError, AttributeError):
315                raise Exception("rig type '" + self.metarig_type + "' has no sample.")
316            else:
317                create_sample(context.active_object)
318            finally:
319                bpy.ops.object.mode_set(mode='EDIT')
320
321        return {'FINISHED'}
322
323
324class EncodeMetarig(bpy.types.Operator):
325    """ Creates Python code that will generate the selected metarig.
326    """
327    bl_idname = "armature.rigify_encode_metarig"
328    bl_label = "Rigify Encode Metarig"
329    bl_options = {'UNDO'}
330
331    @classmethod
332    def poll(self, context):
333        return context.mode == 'EDIT_ARMATURE'
334
335    def execute(self, context):
336        name = "metarig.py"
337
338        if name in bpy.data.texts:
339            text_block = bpy.data.texts[name]
340            text_block.clear()
341        else:
342            text_block = bpy.data.texts.new(name)
343
344        text = write_metarig(context.active_object, layers=True, func_name="create")
345        text_block.write(text)
346        bpy.ops.object.mode_set(mode='EDIT')
347
348        return {'FINISHED'}
349
350
351class EncodeMetarigSample(bpy.types.Operator):
352    """ Creates Python code that will generate the selected metarig
353        as a sample.
354    """
355    bl_idname = "armature.rigify_encode_metarig_sample"
356    bl_label = "Rigify Encode Metarig Sample"
357    bl_options = {'UNDO'}
358
359    @classmethod
360    def poll(self, context):
361        return context.mode == 'EDIT_ARMATURE'
362
363    def execute(self, context):
364        name = "metarig_sample.py"
365
366        if name in bpy.data.texts:
367            text_block = bpy.data.texts[name]
368            text_block.clear()
369        else:
370            text_block = bpy.data.texts.new(name)
371
372        text = write_metarig(context.active_object, layers=False, func_name="create_sample")
373        text_block.write(text)
374        bpy.ops.object.mode_set(mode='EDIT')
375
376        return {'FINISHED'}
377
378
379class EncodeWidget(bpy.types.Operator):
380    """ Creates Python code that will generate the selected metarig.
381    """
382    bl_idname = "mesh.rigify_encode_mesh_widget"
383    bl_label = "Rigify Encode Widget"
384    bl_options = {'UNDO'}
385
386    @classmethod
387    def poll(self, context):
388        return context.mode == 'EDIT_MESH'
389
390    def execute(self, context):
391        name = "widget.py"
392
393        if name in bpy.data.texts:
394            text_block = bpy.data.texts[name]
395            text_block.clear()
396        else:
397            text_block = bpy.data.texts.new(name)
398
399        text = write_widget(context.active_object)
400        text_block.write(text)
401        bpy.ops.object.mode_set(mode='EDIT')
402
403        return {'FINISHED'}
404
405
406#menu_func = (lambda self, context: self.layout.menu("VIEW3D_MT_armature_metarig_add", icon='OUTLINER_OB_ARMATURE'))
407
408#from bl_ui import space_info  # ensure the menu is loaded first
409
410def register():
411    bpy.utils.register_class(DATA_PT_rigify_layer_names)
412    bpy.utils.register_class(DATA_PT_rigify_buttons)
413    bpy.utils.register_class(BONE_PT_rigify_buttons)
414    bpy.utils.register_class(VIEW3D_PT_tools_rigify_dev)
415    bpy.utils.register_class(LayerInit)
416    bpy.utils.register_class(Generate)
417    bpy.utils.register_class(Sample)
418    bpy.utils.register_class(EncodeMetarig)
419    bpy.utils.register_class(EncodeMetarigSample)
420    bpy.utils.register_class(EncodeWidget)
421    #space_info.VIEW3D_MT_armature_add.append(ui.menu_func)
422
423
424def unregister():
425    bpy.utils.unregister_class(DATA_PT_rigify_layer_names)
426    bpy.utils.unregister_class(DATA_PT_rigify_buttons)
427    bpy.utils.unregister_class(BONE_PT_rigify_buttons)
428    bpy.utils.unregister_class(VIEW3D_PT_tools_rigify_dev)
429    bpy.utils.unregister_class(LayerInit)
430    bpy.utils.unregister_class(Generate)
431    bpy.utils.unregister_class(Sample)
432    bpy.utils.unregister_class(EncodeMetarig)
433    bpy.utils.unregister_class(EncodeMetarigSample)
434    bpy.utils.unregister_class(EncodeWidget)
435