gp_render/fn.py

import bpy
import re
from mathutils import Vector
from pathlib import Path

from collections import defaultdict


def create_node(type, tree=None, **kargs):
    '''Get a type, a tree to add in, and optionnaly multiple attribute to set
    return created node
    '''
    tree = tree or bpy.context.scene.node_tree

    node = tree.nodes.new(type)
    for k,v in kargs.items():
        setattr(node, k, v)
    
    return node

def new_aa_node(tree):
    '''create AA node'''
    aa = create_node('CompositorNodeAntiAliasing', tree) # type = ANTIALIASING
    aa.threshold = 0.5
    aa.contrast_limit = 0.5
    aa.corner_rounding = 0.25
    aa.hide = True
    return aa


def set_settings(scene=None):
    if not scene:
        scene = bpy.context.scene
    # specify scene settings for these kind of render
    scene = bpy.context.scene
    scene.eevee.taa_render_samples = 1
    scene.grease_pencil_settings.antialias_threshold = 0
    scene.render.film_transparent = True
    scene.view_settings.view_transform = 'Standard'


def get_render_scene():
    '''Get / Create a scene named Render'''
    render = bpy.data.scenes.get('Render')
    if not render:
        render = bpy.data.scenes.new('Render')
        
        ## link cameras (and lights ?)
        for ob in bpy.context.scene.objects:
            if ob.type in ('CAMERA', 'LIGHT'):
                render.collection.objects.link(ob)

        # set adapted render settings (no AA)
        set_settings(render)
    render.use_nodes = True
    return render

def get_view_layer(name, scene=None):
    '''get viewlayer name
    return existing/created viewlayer
    '''
    if not scene:
        scene = get_render_scene()
    ### pass double letter prefix as suffix
    ## pass_name = re.sub(r'^([A-Z]{2})(_)(.*)', r'\3\2\1', 'name')
    ## pass_name = f'{name}_{passe}'
    pass_vl = scene.view_layers.get(name)
    if not pass_vl:
        pass_vl = scene.view_layers.new(name)
    return pass_vl

### node location tweaks

def real_loc(n):
    if not n.parent:
        return n.location
    return n.location + real_loc(n.parent)

def get_frame_transform(f, node_tree=None):
    '''Return real transform location of a frame node
    only works with one level of nesting (not recursive)
    '''
    if not node_tree:
        node_tree = f.id_data
    if f.type != 'FRAME':
        return
        # return real_loc(f), f.dimensions

    childs = [n for n in node_tree.nodes if n.parent == f]
    # real_locs = [f.location + n.location for n in childs]
    
    xs = [n.location.x for n in childs] + [n.location.x + n.dimensions.x for n in childs]
    ys = [n.location.y for n in childs] + [n.location.y - n.dimensions.y for n in childs]
    xs.sort(key=lambda loc: loc) # x val : ascending 
    ys.sort(key=lambda loc: loc) # ascending # , reversed=True) # y val : descending
    
    loc = Vector((min(xs), max(ys)))
    dim = Vector((max(xs) - min(xs) + 60, max(ys) - min(ys) + 60))
    
    return loc, dim


## get all frames with their real transform.

def bbox(f, frames):
    xs=[]
    ys=[]
    for n in frames[f]: # nodes of passed frame
        # Better as Vectors ?
        if n.type == 'FRAME':
            if n not in frames.keys():
                # print(f'frame {n.name} not in frame list')
                continue
            all_xs, all_ys = bbox(n, frames) # frames[n]
            xs += all_xs
            ys += all_ys

        else:
            loc = real_loc(n)
            xs += [loc.x, loc.x + n.dimensions.x] # + (n.dimensions.x/get_dpi_factor())
            ys += [loc.y, loc.y - n.dimensions.y] # - (n.dimensions.y/get_dpi_factor())


    # margin ~= 30 
    # return xs and ys
    return [min(xs)-30, max(xs)+30], [min(ys)-30, max(ys)+30]

def get_frames_bbox(node_tree):
    '''Return a dic with all frames
    ex: {frame_node: (location, dimension), ...}
    '''

    # create dic of frame object with his direct child nodes nodes
    frames = defaultdict(list)
    frames_bbox = {}
    for n in node_tree.nodes:
        if not n.parent:
            continue
        # also contains frames
        frames[n.parent].append(n)
    
    # Dic for bbox coord
    for f, nodes in frames.items():
        if f.parent:
            continue
        
        xs, ys = bbox(f, frames)
        # xs, ys = bbox(nodes, frames)
        
        ## returning: list of corner coords
        # coords = [
        #     Vector((xs[0], ys[1])),
        #     Vector((xs[1], ys[1])),
        #     Vector((xs[1], ys[0])),
        #     Vector((xs[0], ys[0])),
        # ]
        # frames_bbox[f] = coords

        ## returning: (loc vector, dimensions vector)
        frames_bbox[f] = Vector((xs[0], ys[1])), Vector((xs[1] - xs[0], ys[1] - ys[0]))

    return frames_bbox


## nodes helper functions

def clear_nodegroup(name, full_clear=False):
    '''remove duplication of a nodegroup (.???)
    also remove the base one if full_clear True
    '''

    for ng in reversed(bpy.data.node_groups):
        pattern = name + r'\.\d{3}'

        if not full_clear and ng.users:
            continue

        if re.search(pattern, ng.name):
            bpy.data.node_groups.remove(ng)

        if full_clear and ng.name == name:
            # if full clear
            bpy.data.node_groups.remove(ng)


def rearrange_frames(node_tree):
    frame_d = get_frames_bbox(node_tree) # dic : {frame_node:(loc vector, dimensions vector), ...}
    if not frame_d:
        print('no frame found')
        return

    print([f.name for f in frame_d.keys()])
    ## order the dict by frame.y location
    frame_d = {key: value for key, value in sorted(frame_d.items(), key=lambda pair: pair[1][0].y - pair[1][1].y, reverse=True)}
    frames = [[f, v[0], v[1].y] for f, v in frame_d.items()] # [frame_node, real_loc, real dimensions]

    top = frames[0][1].y # upper node location.y
    offset = 0
    for f in frames:
        ## f[1] : real loc Vector
        ## f[0] : frame

        ## move frame by offset needed (delta between real_loc and "fake" loc , minus offset)
        f[0].location.y = (f[1].y - f[0].location.y) - top - offset
        # f[0].location.y = f[1].y - top - offset
        offset += f[2] + 300 # gap
        
        f[0].update()

def reorder_inputs(ng):
    rl_nodes = [s.links[0].from_node for s in ng.inputs if s.is_linked and s.links and s.links[0].from_node.type == 'R_LAYERS']
    rl_nodes.sort(key=lambda x: x.location.y, reverse=True)
    names = [n.layer for n in rl_nodes]
    inputs_names = [s.name for s in ng.inputs]
    filtered_names = [n for n in names if n in inputs_names]

    for dest, name in enumerate(filtered_names):
        ## rebuild list at each iteration so index are good
        inputs_names = [s.name for s in ng.inputs]
        src = inputs_names.index(name)
        # reorder on node_tree not directly on node!
        ng.node_tree.inputs.move(src, dest)

def reorder_outputs(ng):
    ordered_out_name = [nis.name for nis in ng.inputs if nis.name in [o.name for o in ng.outputs]]
    for s_name in ordered_out_name:
        all_outnames = [o.name for o in ng.outputs]
        # reorder on nodetree, not on node !
        ng.node_tree.outputs.move(all_outnames.index(s_name), ordered_out_name.index(s_name))

def clear_disconnected(fo):
    for inp in reversed(fo.inputs):
        if not inp.is_linked:
            print(f'Deleting unlinked fileout slot: {inp.name}')
            fo.inputs.remove(inp)

def reorder_fileout(fo, ng=None):
    if not ng: # get connected nodegroup
        for s in fo.inputs:
            if s.is_linked and s.links and s.links[0].from_node.type == 'GROUP':
                ng = s.links[0].from_node
                break
    if not ng:
        print(f'No nodegroup to refer to filter {fo.name}')
        return
    ordered = [o.links[0].to_socket.name for o in ng.outputs if o.is_linked and o.is_linked and o.links[0].to_node == fo]
    for s_name in ordered:
        all_outnames = [s.name for s in fo.inputs] # same as [fs.path for fs in fo.file_slots]
        fo.inputs.move(all_outnames.index(s_name), ordered.index(s_name))

def reorganise_NG_nodegroup(ng):
    '''refit node content to avoid overlap'''
    ngroup = ng.node_tree
    ng_in = ngroup.nodes.get('Group Input')
    offset = 35
    y = 0
    for s in ng_in.outputs:
        if s.is_linked:
            s.links[0].to_node.location.y = y
        y -= offset

def connect_to_group_output(n):
    for o in n.outputs:
        if o.is_linked:
            if o.links[0].to_node.type == 'GROUP_OUTPUT':
                return o.links[0].to_socket
            val = connect_to_group_output(o.links[0].to_node)
            if val:
                return val
    return False

def connect_to_group_input(n):
    for i in n.inputs:
        if i.is_linked:
            if i.links[0].from_node.type == 'GROUP_INPUT':
                return i.links[0].from_socket
            val = connect_to_group_input(i.links[0].from_node)
            if val:
                return val
    return False


def clear_nodegroup_content_if_disconnected(ngroup):
    '''Get a nodegroup.node_tree
    delete orphan nodes that are not connected from group input node
    '''
    for n in reversed(ngroup.nodes):
        if n.type in ('GROUP_INPUT', 'GROUP_OUTPUT'):
            continue
        if not connect_to_group_input(n) and not connect_to_group_output(n): # is disconnected from both side
            ngroup.nodes.remove(n)

def clean_nodegroup_inputs(ng):
    '''Clear inputs to output of passed nodegroup if not connected'''
    ngroup = ng.node_tree
    for i in range(len(ng.inputs))[::-1]:
        print(i)
        if not ng.inputs[i].is_linked: # and not any(x.layer == s.name for x in rl_nodes)
            ngroup.inputs.remove(ngroup.inputs[i])
    
    # clear_nodegroup_content_if_disconnected(ngroup)

def random_color(alpha=False):
    import random
    if alpha:
        return (random.uniform(0,1), random.uniform(0,1), random.uniform(0,1), 1)
    return (random.uniform(0,1), random.uniform(0,1), random.uniform(0,1))


def get_numbered_output(out, slot_name):
    '''Return output slot name without looking for numbering ???_
    '''
    pattern = r'^(?:\d{3}_)?' # optional non capture group of 3 digits + _
    pattern = f'{pattern}{slot_name}'
    for inp in out.inputs:
        if re.match(pattern, inp.name):
            return inp


def add_fileslot_number(fs, number):
    elems = fs.path.split('/')
    for i, e in enumerate(elems):
        if re.match(r'^\d{3}_', e):
            elems[i] = re.sub(r'^(\d{3})', lambda x: str(number).zfill(3), e)
        else:
            elems[i] = f'{str(number).zfill(3)}_{e}'
    new = '/'.join(elems)
    fs.path = new
    return new

def renumber(fo, offset=10):
    '''Force renumber all the slots with a 3'''

    if fo.type != 'OUTPUT_FILE': return
    ct = 0
    for fs in fo.file_slots:
        add_fileslot_number(fs, ct)
        ct += offset

def get_num(string):
    num = re.search(r'^(\d{3})_', string)
    if num:
        return int(num.group(1))
""" 
def renumber_keep(fo, offset=10):
    '''Force renumber all the slots with a 3'''

    if fo.type != 'OUTPUT_FILE': return

    renum = re.compile(r'^(\d{3})_')
    ct = 0

    last_idx = len(fo.file_slots) -1
    prev = None
    prev_num = None
    for idx, fs in enumerate(fo.file_slots):
        # if idx!=last_idx:
        #     fs_next = fo.file_slots[idx+1]
        if idx == last_idx: # handle last
            if idx > 0:
                prev = fo.file_slots[idx-1]
                num = renum.search(prev)
                if not num:
                    add_fileslot_number(fs, ct)
                else:
                    add_fileslot_number(fs, int(num.group(1)) + offset)
            else:
                add_fileslot_number(fs, 0) # there is only one slot (maybe don't number ?)
            break

        # update the ct with the current taken number if any
        number = renum.search(fs.path)
        if number:
            prev = fs
            ct = number + offset
            continue # skip already numbered

        # analyse all next slots until there is numbered
        divider = None
        for i in range(idx + 1, len(fo.file_slots) - idx):
            next_num = renum.search(fo.file_slots[idx + i])
            if next_num:
                divider = i-1
                break

        if idx == 0:
            prev_num = 0
            prev = None
        else:
            prev = fo.file_slots[idx-1]
            prev_num = renum.search(prev.path)
            if prev_num:
                prev_num = prev_num.group(1)
        
        if not divider: # just use prev and next if prev
            if not prev:
                add_fileslot_number(fs, ct)            

        # first check if it has a number (if not bas)
        prev = fs
        ct += offset
"""

def delete_numbering(fo): # padding=3
    '''Delete prefix numbering on all slots on passed file output'''

    if fo.type != 'OUTPUT_FILE': return
    for fs in fo.file_slots:
        elems = fs.path.split('/')
        for i, e in enumerate(elems):
            elems[i] = re.sub(r'^\d{3}_', '', e)
        
        new = '/'.join(elems)
        fs.path = new
        


def nodegroup_merge_inputs(ngroup):
    '''Get a nodegroup
    merge every group inputs with alpha over
    then connect to antialias and a new output
    '''

    ng_in = ngroup.nodes.get('Group Input')
    ng_out = ngroup.nodes.get('Group Output')

    x, y = ng_in.location.x + 200, 0

    offset_x, offset_y = 150, -100

    # merge all inputs in alphaover nodes
    prev = None
    for i in range(len(ng_in.outputs)-1): # skip waiting point
        inp = ng_in.outputs[i]
        if not prev:
            prev = ng_in
            continue

        # live connect
        ao = create_node('CompositorNodeAlphaOver', tree=ngroup, location=(x,y), hide=True)
        ngroup.links.new(prev.outputs[0], ao.inputs[1])
        ngroup.links.new(inp, ao.inputs[2])
        
        x += offset_x
        y += offset_y
        prev = ao

    ## create a merged name as output ??
    aa = new_aa_node(ngroup)
    aa.location = (ao.location.x + 200, ao.location.y)
    ngroup.links.new(ao.outputs[0], aa.inputs[0]) # node_tree
    
    # create one input and link
    out = ngroup.outputs.new('NodeSocketColor', ngroup.inputs[0].name)
    ngroup.links.new(aa.outputs[0], ng_out.inputs[0])