rig_picker/op_shape.py

import bpy
from .func_shape import get_picker_datas
from .utils import is_shape, find_mirror, link_mat_to_object
#import os
from pathlib import Path
import json


class RP_OT_create_shape(bpy.types.Operator):
    bl_label = 'Create UI shape'
    bl_idname = 'rigpicker.create_shape'
    #bl_options = {'REGISTER', 'UNDO'}

    @classmethod
    def poll(cls, context):
        return (context.object and context.object.mode == 'POSE')

    def execute(self,context):
        scn = context.scene
        vl = context.view_layer

        offset = 0
        for bone in context.selected_pose_bones:
            name = bone.name

            mesh = bpy.data.meshes.new(name)
            ob = bpy.data.objects.new(name, mesh)

            ob.rig_picker.name = name

            verts=[(0,0,0)]
            edges = []
            faces =[]

            mesh.from_pydata(verts, edges, faces)

            for c in scn.rig_picker.canvas.users_collection:
                c.objects.link(ob)

            ob.location.z = 0.05
            ob.location.x = offset

            for obj in scn.objects:
                obj.select_set(False)

            vl.objects.active = ob
            ob.select_set(True)

            offset += 0.05

        return {'FINISHED'}


class RP_OT_name_from_bone(bpy.types.Operator):
    bl_label = 'Name Shape from selected bones'
    bl_idname = 'rigpicker.name_from_bone'
    #bl_options = {'REGISTER', 'UNDO'}

    def execute(self,context):
        scene = context.scene
        rig = scene.rig_picker.rig
        bone = rig.data.bones.active

        if bone:
            context.object.name = bone.name
            context.object.rig_picker.name = bone.name

        return {'FINISHED'}


class RP_OT_mirror_shape(bpy.types.Operator):
    bl_label = 'Mirror UI shape'
    bl_idname = 'rigpicker.mirror_shape'
    #bl_options = {'REGISTER', 'UNDO'}

    def execute(self,context):
        scn = context.scene

        for ob in bpy.context.selected_objects:

            name = find_mirror(ob.name)
            link_mat_to_object(ob)

            if not name:
                name = ob.name + '_flip'

            old_shape = bpy.data.objects.get(name)
            old_mat = None
            if old_shape:
                old_mat = next((sl.material for sl in old_shape.material_slots if sl.material), None)
                bpy.data.objects.remove(old_shape)

            mirror_ob = ob.copy()
            mirror_ob.data = ob.data
            mirror_ob.name = name

            if old_mat:
                #mirror_ob.data.materials.clear()
                #mirror_ob.data.materials.append(None)
                #mirror_ob.material_slots[0].link = 'OBJECT'

                mirror_ob.material_slots[0].material = old_mat
            #mirror_ob = bpy.data.objects.new(name,ob.data.copy())

            for c in ob.users_collection:
                c.objects.link(mirror_ob)

            if scn.rig_picker.symmetry:
                symmetry_loc = scn.rig_picker.symmetry.matrix_world.to_translation()[0]
            else:
                symmetry_loc = 0

            #print(symmetry_loc)
            mirror_ob.matrix_world = ob.matrix_world

            #if mirror_ob.location[0] < symmetry_loc:
            mirror_ob.location.x = symmetry_loc + (symmetry_loc - ob.location.x)
            #else:
            #    mirror_ob.location[0]= symmetry_loc+ (symmetry_loc- ob.location[0])

            mirror_ob.rotation_euler.y = -ob.rotation_euler.y
            mirror_ob.rotation_euler.z = -ob.rotation_euler.z

            mirror_ob.scale.x = -ob.scale.x

            for key, value in ob.items():
                if key not in ['_RNA_UI','cycles']:
                    mirror_ob[key] = value

            if ob.rig_picker.shape_type == 'BONE':
                mirror_ob.rig_picker.name = find_mirror(ob.rig_picker.name)

            elif ob.rig_picker.shape_type == 'FUNCTION':
                args = {}
                for key,value in eval(ob.rig_picker.arguments).items():
                    if type(value) == list:
                        mirrored_value = []
                        for item in value:
                            mirrored_value.append(find_mirror(item))

                    elif type(value) == str:
                        mirrored_value = find_mirror(value)
                    args[key] = mirrored_value

                mirror_ob.rig_picker.arguments = str(args)

        return {'FINISHED'}


class RP_OT_select_shape_type(bpy.types.Operator):
    bl_label = 'Select Shape by Type'
    bl_idname = 'rigpicker.select_shape_type'
    #bl_options = {'REGISTER', 'UNDO'}

    shape_type = bpy.props.EnumProperty(items =[(i.upper(), i, '') for i in ('Bone', 'Display', 'Function')])

    def draw(self,context):
        layout = self.layout

        col = layout.column()
        col.prop(self,'shape_type', expand=True)

    def execute(self,context):
        scene = context.scene
        #print(self.shape_type)
        canvas = context.scene.rig_picker.canvas
        if canvas:
            for ob in [o for o in bpy.data.objects if o.layers==canvas.layers]:
                if ob.type in ['MESH', 'CURVE', 'FONT'] and ob.rig_picker.shape_type == self.shape_type:
                    ob.select = True

        return {'FINISHED'}

    def invoke(self, context, event):
        wm = context.window_manager
        return wm.invoke_props_dialog(self, width=150)


class RP_OT_save_picker(bpy.types.Operator):
    bl_label = 'Store UI Data'
    bl_idname = 'rigpicker.save_picker'

    def execute(self, context):
        scn = context.scene
        canvas= scn.rig_picker.canvas
        rig = scn.rig_picker.rig
        shapes = [o for o in scn.objects if o != canvas and is_shape(o)]

        if not rig:
            self.report({'ERROR'}, 'Choose a Rig')
            return {'CANCELLED'}

        if not canvas:
            self.report({'ERROR'}, 'Choose a Canvas')
            return {'CANCELLED'}


        data = get_picker_datas(shapes, canvas, rig)

        picker_path = Path(bpy.path.abspath(scn.rig_picker.destination))
        picker_path.write_text(json.dumps(data))

        bpy.ops.rigpicker.reload_picker()


        return {'FINISHED'}

classes = (
    RP_OT_create_shape,
    RP_OT_name_from_bone,
    RP_OT_mirror_shape,
    RP_OT_select_shape_type,
    RP_OT_save_picker
)

register, unregister = bpy.utils.register_classes_factory(classes)