rig_picker/draw_handlers.py

import bpy
import gpu
from gpu_extras.batch import batch_for_shader
import blf
from pathlib import Path
from .constants import PICKERS
from .core.picker import Picker, PickerGroup
import json


def draw_rect_2d(position, width, height, color):
    """
    Draw a 2d rectangele.

    :arg position: Position of the lower left corner.
    :type position: 2D Vector
    :arg width: Width of the rect.
    :type width: float
    :arg height: Height of the rect.
    :type height: float
    """
    
    coords = ((0, 0), (1, 0), (1, 1), (0, 1))

    shader = gpu.shader.from_builtin('UNIFORM_COLOR')
    batch = batch_for_shader(
        shader, 'TRI_FAN',
        {"pos": coords},
    )

    with gpu.matrix.push_pop():
        gpu.matrix.translate(position)
        gpu.matrix.scale((width, height))
        shader.uniform_float("color", color)

        batch.draw(shader)


def draw_callback_view():
    sp = bpy.context.space_data

    if not sp or not sp.tree_type == 'RigPickerTree':
        return

    if not sp.pin:
        sp.pin = True
        sp.show_region_ui = False

    ob = bpy.context.object
    if not ob or ob.type !='ARMATURE' or not ob.data.rig_picker.sources:
        return

    if ob not in PICKERS:
        if 'pickers' in ob.data.rig_picker:
            picker_datas = [s.to_dict() for s in ob.data.rig_picker['pickers']]
        else:
            picker_datas = []
            for picker in ob.data.rig_picker.sources:
                picker_path = Path(bpy.path.abspath(picker.source, library=ob.data.library))

                if not picker_path.exists():
                    print(f'Picker path not exists: {picker_path.resolve()}')
                    continue

                print('Load picker from', picker_path.resolve())
                picker_data = json.loads(picker_path.read_text(encoding='utf-8'))
                picker_datas.append(picker_data)
                #shapes = [s.to_dict() for s in ob.data.rig_picker['shapes']]

        PICKERS[ob] = PickerGroup(ob, picker_datas)#[Picker(ob, shapes=picker_data) for picker_data in picker_datas]

        # if 'picker' in ob.data.rig_picker:
        #     picker_datas = [s.to_dict() for s in ob.data.rig_picker['picker']]
        # else:
        #     picker_path = Path(bpy.path.abspath(ob.data.rig_picker.source, library=ob.data.library))

        #     if not picker_path.exists():
        #         print(f'Picker path not exists: {picker_path.resolve()}')
        #         return

        #     print('Load picker from', picker_path.resolve())
        #     picker_datas = json.loads(picker_path.read_text())
        #     #shapes = [s.to_dict() for s in ob.data.rig_picker['shapes']]
        # PICKERS[ob] = Picker(ob, shapes=picker_datas)

    picker_group = PICKERS.get(ob)   
    picker_group.draw()

def draw_callback_px():
    sp = bpy.context.space_data
    if not sp.tree_type == 'RigPickerTree':
        return

    ob = bpy.context.object

    picker_group = PICKERS.get(ob)
    if not picker_group or not picker_group.tooltip:
        return

    text = picker_group.tooltip
    mouse = picker_group.tooltip_mouse
    ui_scale = bpy.context.preferences.system.ui_scale

    #print('Draw text', text)
    font_id = 0
    blf.size(font_id, int(13 * ui_scale))

    text_size = blf.dimensions(font_id, text)
    text_pos = (mouse[0] - text_size[0]*0.33, mouse[1] - (34*ui_scale))

    bg_margins = [8, 4]
    bg_pos = (text_pos[0] - bg_margins[0], text_pos[1] - bg_margins[1]-1)
    bg_size = (text_size[0] + 2*bg_margins[0], text_size[1] + 2*bg_margins[1])

    gpu.state.blend_set('ALPHA')
    draw_rect_2d(bg_pos, *bg_size, (0, 0, 0, 0.75))
    gpu.state.blend_set('NONE')


    #blf.dimensions(font_id, text)
    blf.enable(font_id, blf.SHADOW)


    # BLF drawing routine
    blf.position(font_id, round(text_pos[0]), round(text_pos[1]), 0)
    blf.color(font_id, 1, 1, 1, 1)

    blf.shadow(font_id , 5, 0.0, 0.0, 0.0, 1)
    blf.shadow_offset(font_id, 2, -2)

    blf.draw(font_id, text)

    blf.disable(font_id, blf.SHADOW)

handle_view = None
handle_pixel = None


def register():
    global handle_view
    global handle_pixel

    handle_view = bpy.types.SpaceNodeEditor.draw_handler_add(draw_callback_view, (), 'WINDOW', 'POST_VIEW')
    handle_pixel = bpy.types.SpaceNodeEditor.draw_handler_add(draw_callback_px, (), 'WINDOW', 'POST_PIXEL')

def unregister():
    global handle_view
    global handle_pixel

    bpy.types.SpaceNodeEditor.draw_handler_remove(handle_view, 'WINDOW')
    bpy.types.SpaceNodeEditor.draw_handler_remove(handle_pixel, 'WINDOW')

    handle_view = None
    handle_pixel = None