bone_widget/shape_utils.py

import bpy
from .transform_utils import get_bone_matrix, transform_matrix
from mathutils import Matrix
#from bone_widget import ctx


def rename_shape(shape, bone, prefix=''):
    shape.name = shape.data.name = prefix + bone.name
    #print('Rename Shape', bone, shape.name)

def get_bone_from_shape(shape):
    armatures = [o for o in bpy.context.scene.objects if o.type == 'ARMATURE']
    return next((b for a in armatures for b in a.pose.bones if b.custom_shape is shape), None)

'''
def set_shape(bone, shape):
    shape.matrix_world = get_bone_matrix(bone)
    bone.custom_shape = shape
    ctx.rename_shape(shape, bone)
    ctx.link_to_widget_col(shape)
'''

def get_side(name, fallback=None):
    name = name.lower()
    if name.endswith(('.l', '_l')) or name.startswith(('l.', 'l_')):
        return 'LEFT'
    elif name.endswith(('.r', '_r')) or name.startswith(('r.', 'r_')):
        return 'RIGHT'

    return fallback

def get_flipped_name(name):
    import re

    def flip(match, start=False):
        if not match.group(1) or not match.group(2):
            return
        
        sides = {
            'R' : 'L',
            'r' : 'l',
            'L' : 'R',
            'l' : 'r',
        }

        if start:
            side, sep = match.groups()
            return sides[side] + sep
        else:
            sep, side, num = match.groups()
            return sep + sides[side] + (num or '')

    start_reg = re.compile(r'(l|r)([._-])', flags=re.I)

    if start_reg.match(name):
        flipped_name = start_reg.sub(lambda x: flip(x, True), name)
    else:
        flipped_name = re.sub(r'([._-])(l|r)(\.\d+)?$', flip, name, flags=re.I)

    return flipped_name

def get_flipped_bone(bone):
    flipped_name = get_flipped_name(bone.name)

    if flipped_name == bone.name:
        return
    
    return bone.id_data.pose.bones.get(flipped_name)

def link_to_col(shape, col):
    for c in shape.users_collection:
        c.objects.unlink(shape)

    col.objects.link(shape)

def custom_shape_matrix(bone):
    loc = bone.custom_shape_translation.copy()
    rot = bone.custom_shape_rotation_euler.copy()
    scale = bone.custom_shape_scale_xyz.copy()

    if bone.use_custom_shape_bone_size:
        loc /= bone.bone.length
    else:
        scale /= bone.bone.length

    return Matrix.LocRotScale(loc, rot, scale)

def apply_custom_shape_transform(bone):
    shape = bone.custom_shape

    if not shape:
        return

    if bpy.app.version_string < '3.0.0':
        scale = bone.custom_shape_scale

        if round(scale, 5) == 0:
            return
        
        if not bone.use_custom_shape_bone_size:
            scale /= bone.bone.length

        mat = transform_matrix(scale=(1/scale,)*3)

        shape.data.transform(mat)
        bone.custom_shape_scale = 1

    else:
        mat = custom_shape_matrix(bone)

        scale = (1, 1, 1)
        if round(mat.to_scale().length, 5) == 0:
            loc, rot, scale = mat.decompose()
            mat = Matrix.LocRotScale(loc, rot, (1, 1, 1))
            scale = bone.custom_shape_scale_xyz 

        shape.data.transform(mat)

        bone.custom_shape_translation = (0, 0, 0)
        bone.custom_shape_rotation_euler = (0, 0, 0)
        bone.custom_shape_scale_xyz = scale
    
def get_clean_shape(bone, shape, separate=True, rename=True, 
        col=None, match=True, prefix='', apply_transforms=True):

    #bone.custom_shape = None

    #Remove bone shape if no other bone are using it
    if bone.custom_shape:
        old_bones = get_bones(bone.custom_shape)
        old_bones.remove(bone)

        print('old_bones', old_bones)

        if old_bones:
            rename_shape(bone.custom_shape, old_bones[0], prefix=prefix)
        elif bone.custom_shape != shape:
            bpy.data.objects.remove(bone.custom_shape)

    if separate:
        shape = shape.copy()
        shape.data = shape.data.copy()

    bone.custom_shape = shape
    
    if apply_transforms:
        apply_custom_shape_transform(bone)
        bone.use_custom_shape_bone_size = True

    if match:
        shape.matrix_world = get_bone_matrix(bone)

    if rename:
        rename_shape(shape, bone, prefix=prefix)

    if col:
        link_to_col(shape, col)

    return shape


def get_bones(shape):
    armatures = [o for o in bpy.context.scene.objects if o.type == 'ARMATURE']
    return [b for a in armatures for b in a.pose.bones if b.custom_shape == shape]

def get_bone(shape):
    bones = get_bones(shape)
    if bones:
        return bones[0]

def symmetrize_bone_shape(bone, prefix=None):
    active_side = get_side(bone.name, 'LEFT')

    shape = bone.custom_shape
    flipped_bone = get_flipped_bone(bone)

    if not flipped_bone or not shape:
        return

    if bpy.app.version_string < '3.0.0':
        flipped_bone.custom_shape_scale = b.custom_shape_scale
    else:
        flipped_bone.custom_shape_translation = bone.custom_shape_translation
        flipped_bone.custom_shape_rotation_euler = bone.custom_shape_rotation_euler
        flipped_bone.custom_shape_scale_xyz = bone.custom_shape_scale_xyz

    flipped_shape = flipped_bone.custom_shape

    if flipped_shape:
        flipped_bone.custom_shape = None

        old_bone = get_bone(flipped_shape)
        if old_bone:
            rename_shape(flipped_shape, old_bone, prefix=prefix)
        else:
            bpy.data.objects.remove(flipped_shape)

    flipped_shape = shape.copy()
    flipped_shape.data = flipped_shape.data.copy()
    flipped_shape.data.transform(transform_matrix(scale=(-1, 1, 1)))

    flipped_bone.custom_shape = get_clean_shape(flipped_bone, flipped_shape, rename=True, separate=False, prefix=prefix, apply_transforms=False)


    flipped_bone.use_custom_shape_bone_size = bone.use_custom_shape_bone_size

    for c in shape.users_collection:
        c.objects.link(flipped_bone.custom_shape)
    
    return flipped_bone.custom_shape