asset_library/core/bl_utils.py

"""
Generic Blender functions
"""
import json
from pathlib import Path
from fnmatch import fnmatch
from typing import Any, List, Iterable, Optional, Tuple
Datablock = Any

import bpy
from bpy_extras import asset_utils
#from asset_library.constants import RESOURCES_DIR
#from asset_library.common.file_utils import no
from os.path import abspath
import subprocess

from .file_utils import norm_str


class attr_set():
    '''Receive a list of tuple [(data_path, "attribute" [, wanted value)] ]
    entering with-statement : Store existing values, assign wanted value (if any)
    exiting with-statement: Restore values to their old values
    '''

    def __init__(self, attrib_list):
        self.store = []
        # item = (prop, attr, [new_val])
        for item in attrib_list:
            prop, attr = item[:2]
            self.store.append( (prop, attr, getattr(prop, attr)) )
        
        for item in attrib_list:
            prop, attr = item[:2]

            if len(item) >= 3:
                try:
                    setattr(prop, attr, item[2])
                except TypeError:
                    print(f'Cannot set attribute {attr} to {prop}')

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, exc_traceback):
        self.restore()

    def restore(self):
        for prop, attr, old_val in self.store:
            setattr(prop, attr, old_val)


def unique_name(name, names):
    if name not in names:
        return name

    i = 1
    org_name = name
    while name in names:
        name = f'{org_name}.{i:03d}'
        i += 1
    
    return name


def get_overriden_col(ob, scene=None):
    scn = scene or bpy.context.scene

    cols = [c for c in bpy.data.collections if scn.user_of_id(c)]
    
    return next((c for c in cols if ob in c.all_objects[:] 
           if all(not c.override_library for c in get_col_parents(c))), None)


def get_view3d_persp():
    windows = bpy.context.window_manager.windows
    view_3ds = [a for w in windows for a in w.screen.areas if a.type == 'VIEW_3D']
    view_3d = next((a for a in view_3ds if a.spaces.active.region_3d.view_perspective == 'PERSP'), view_3ds[0])
    return view_3d


def get_viewport():
    screen = bpy.context.screen

    areas = [a for a in screen.areas if a.type == 'VIEW_3D']
    if not areas:
        return 

    areas.sort(key=lambda x : x.width*x.height)

    return areas[-1]


def biggest_asset_browser_area(screen: bpy.types.Screen) -> Optional[bpy.types.Area]:
    """Return the asset browser Area that's largest on screen.

    :param screen: context.window.screen

    :return: the Area, or None if no Asset Browser area exists.
    """

    def area_sorting_key(area: bpy.types.Area) -> Tuple[bool, int]:
        """Return area size in pixels."""
        return (area.width * area.height)

    areas = list(suitable_areas(screen))
    if not areas:
        return None

    return max(areas, key=area_sorting_key)


def suitable_areas(screen: bpy.types.Screen) -> Iterable[bpy.types.Area]:
    """Generator, yield Asset Browser areas."""

    for area in screen.areas:
        space_data = area.spaces[0]
        if not asset_utils.SpaceAssetInfo.is_asset_browser(space_data):
            continue
        yield area


def area_from_context(context: bpy.types.Context) -> Optional[bpy.types.Area]:
    """Return an Asset Browser suitable for the given category.

    Prefers the current Asset Browser if available, otherwise the biggest.
    """

    space_data = context.space_data
    if asset_utils.SpaceAssetInfo.is_asset_browser(space_data):
        return context.area

    # Try the current screen first.
    browser_area = biggest_asset_browser_area(context.screen)
    if browser_area:
        return browser_area

    for win in context.window_manager.windows:
        if win.screen == context.screen:
            continue
        browser_area = biggest_asset_browser_area(win.screen)
        if browser_area:
            return browser_area

    return None


def activate_asset(
    asset: bpy.types.Action, asset_browser: bpy.types.Area, *, deferred: bool
) -> None:
    """Select & focus the asset in the browser."""

    space_data = asset_browser.spaces[0]
    assert asset_utils.SpaceAssetInfo.is_asset_browser(space_data)
    space_data.activate_asset_by_id(asset, deferred=deferred)


def active_catalog_id(asset_browser: bpy.types.Area) -> str:
    """Return the ID of the catalog shown in the asset browser."""
    return params(asset_browser).catalog_id


def get_asset_space_params(asset_browser: bpy.types.Area) -> bpy.types.FileAssetSelectParams:
    """Return the asset browser parameters given its Area."""
    space_data = asset_browser.spaces[0]
    assert asset_utils.SpaceAssetInfo.is_asset_browser(space_data)
    return space_data.params


def refresh_asset_browsers():
    for area in suitable_areas(bpy.context.screen):
        bpy.ops.asset.library_refresh({"area": area, 'region': area.regions[3]})


def tag_redraw(screen: bpy.types.Screen) -> None:
    """Tag all asset browsers for redrawing."""

    for area in suitable_areas(screen):
        area.tag_redraw()

# def get_blender_command(file=None, script=None, background=True, **args):
#     '''Return a Blender Command as a list to be used in a subprocess'''

#     cmd = [bpy.app.binary_path]

#     if file:
#         cmd += [str(file)]
#     if background:
#         cmd += ['--background']
#     if script:
#         cmd += ['--python', str(script)]
#     if args:
#         cmd += ['--']
#         for k, v in args.items():
#             cmd += [f"--{k.replace('_', '-')}", str(v)]
            
#     return cmd

def norm_value(value):
    if isinstance(value, (tuple, list)):
        values = []
        for v in value:
            if not isinstance(v, str):
                v = json.dumps(v)
            values.append(v)

        return values       

    if isinstance(value, Path):
        return str(value)

    if not isinstance(value, str):
        value = json.dumps(value)
    return value


def norm_arg(arg_name, format=str.lower, prefix='--', separator='-'):
    arg_name = norm_str(arg_name, format=format, separator=separator)

    return prefix + arg_name


def get_bl_cmd(blender=None, background=False, factory_startup=False, focus=True, blendfile=None, script=None, **kargs):
    cmd = [str(blender)] if blender else [bpy.app.binary_path]

    if background:
        cmd += ['--background']

    if not focus and not background:
        cmd += ['--no-window-focus']
        cmd += ['--window-geometry', '5000', '0', '10', '10']

    cmd += ['--python-use-system-env']

    if factory_startup:
        cmd += ['--factory-startup']

    if blendfile:
        cmd += [str(blendfile)]
        
    if script:
        cmd += ['--python', str(script)]
    
    if kargs:
        cmd += ['--']     
        for k, v in kargs.items():
            k = norm_arg(k)
            v = norm_value(v)

            cmd += [k]
            if isinstance(v, (tuple, list)):
                cmd += v
            else:
                cmd += [v]

    return cmd


def get_addon_prefs():
    addon_name = __package__.split('.')[0]
    return bpy.context.preferences.addons[addon_name].preferences


def get_col_parents(col, root=None, cols=None):
    '''Return a list of parents collections of passed col
    root : Pass a collection to search in (recursive)
    else search in master collection
    '''
    if cols is None:
        cols = []

    if root == None:
        root = bpy.context.scene.collection

    for sub in root.children:
        if sub == col:
            cols.append(root)

        if len(sub.children):
            cols = get_col_parents(col, root=sub, cols=cols)
    return cols


def get_overriden_col(ob, scene=None):
    '''Get the collection use for making the override'''
    scn = scene or bpy.context.scene

    cols = [c for c in bpy.data.collections if scn.user_of_id(c)]
    
    return next((c for c in cols if ob in c.all_objects[:] 
           if all(not c.override_library for c in get_col_parents(c))), None)


def load_assets_from(filepath: Path) -> List[Datablock]:
    if not has_assets(filepath):
        # Avoid loading any datablocks when there are none marked as asset.
        return []

    # Append everything from the file.
    with bpy.data.libraries.load(str(filepath)) as (
        data_from,
        data_to,
    ):
        for attr in dir(data_to):
            setattr(data_to, attr, getattr(data_from, attr))

    # Iterate over the appended datablocks to find assets.
    def loaded_datablocks() -> Iterable[Datablock]:
        for attr in dir(data_to):
            datablocks = getattr(data_to, attr)
            for datablock in datablocks:
                yield datablock

    loaded_assets = []
    for datablock in loaded_datablocks():
        if not getattr(datablock, "asset_data", None):
            continue

        # Fake User is lost when appending from another file.
        datablock.use_fake_user = True
        loaded_assets.append(datablock)
    return loaded_assets


def has_assets(filepath: Path) -> bool:
    with bpy.data.libraries.load(str(filepath), assets_only=True) as (
        data_from,
        _,
    ):
        for attr in dir(data_from):
            data_names = getattr(data_from, attr)
            if data_names:
                return True
    return False


def copy_frames(start, end, offset, path):
    for i in range (start, end):
        src = path.replace('####', f'{i:04d}')
        dst = src.replace(src.split('_')[-1].split('.')[0], f'{i+offset:04d}')
        shutil.copy2(src, dst)


def split_path(path) :
    try :
        bone_name = path.split('["')[1].split('"]')[0]
    except :
        bone_name = None
    try :
        prop_name = path.split('["')[2].split('"]')[0]
    except :
        prop_name = path.split('.')[-1]

    return bone_name, prop_name


def get_asset_type(asset_type):
    data_types = {  p.fixed_type.identifier: p.identifier for p in 
                bpy.types.BlendData.bl_rna.properties if hasattr(p, 'fixed_type')}
    return data_types[asset_type]


def load_datablocks(src, names=None, type='objects', link=True, expr=None, assets_only=False) -> list:
    return_list = not isinstance(names, str)
    names = names or []

    if not isinstance(names, (list, tuple)):
        names = [names]
    
    if isinstance(expr, str):
        pattern = expr
        expr = lambda x : fnmatch(x, pattern)
    
    with bpy.data.libraries.load(str(src), link=link,assets_only=assets_only) as (data_from, data_to):
        datablocks = getattr(data_from, type)
        if expr:
            names += [i for i in datablocks if expr(i)]
        elif not names:
            names = datablocks
        
        setattr(data_to, type, names)

        datablocks = getattr(data_to, type)

    if return_list:
        return datablocks

    elif datablocks:
        return datablocks[0]

"""
# --- Collection handling
"""

def col_as_asset(col, verbose=False):
    if col is None:
        return
    if verbose:
        print('linking:', col.name)
    pcol = bpy.data.collections.new(col.name)
    bpy.context.scene.collection.children.link(pcol)
    pcol.children.link(col)
    pcol.asset_mark()
    return pcol

def load_col(filepath, name, link=True, override=True, rig_pattern=None, context=None):
    '''Link a collection by name from a file and override if has armature'''

    # with bpy.data.libraries.load(filepath, link=link) as (data_from, data_to):
    #     data_to.collections = [c for c in data_from.collections if c == name]
    # if not data_to.collections:
    #     return
    # return data_to.collections[0]
    context = context or bpy.context

    col = load_datablocks(filepath, name, link=link, type='collections')
    
    ## create instance object
    inst = bpy.data.objects.new(col.name, None)
    inst.instance_collection = col
    inst.instance_type = 'COLLECTION'
    context.scene.collection.objects.link(inst)

    # make active
    inst.select_set(True)
    context.view_layer.objects.active = inst

    ## simple object (no armatures)
    if not link or not override:
        return inst
    if not next((o for o in col.all_objects if o.type == 'ARMATURE'), None):
        return inst
    
    ## Create the override
    # Search
    parent_cols = inst.users_collection
    child_cols = [child for pcol in parent_cols for child in pcol.children]

    params = {'active_object': inst, 'selected_objects': [inst]}
    try:
        bpy.ops.object.make_override_library(params)

        ## check which collection is new in parents collection
        asset_col = next((c for pcol in parent_cols for c in pcol.children if c not in child_cols), None)
        if not asset_col:
            print('Overriden, but no collection found !!')
            return
        
        for ob in asset_col.all_objects:
            if ob.type != 'ARMATURE':
                continue
            if rig_pattern and not fnmatch(ob.name, rig_pattern):
                continue

            ob.hide_select = ob.hide_viewport = False
            ob.select_set(True)
            context.view_layer.objects.active = ob
            print(ob.name)
            return ob

    except Exception as e:
        print(f'Override failed on {col.name}')
        print(e)
    
    return inst


def get_preview(asset_path='', asset_name=''):
    asset_preview_dir = Path(asset_path).parents[1]
    name = asset_name.lower()
    return next((f for f in asset_preview_dir.rglob('*') if f.stem.lower().endswith(name)), None)

def get_object_libraries(ob):
    if ob is None:
        return []
    
    libraries = [ob.library]
    if ob.data:
        libraries += [ob.data.library]

    if ob.type in ('MESH', 'CURVE'):
        libraries += [m.library for m in ob.data.materials if m]
    
    filepaths = []
    for l in libraries:
        if not l or not l.filepath:
            continue
        
        absolute_filepath = abspath(bpy.path.abspath(l.filepath, library=l))
        if absolute_filepath in filepaths:
            continue

        filepaths.append(absolute_filepath)

    return filepaths


def clean_name(name):
    if re.match(r'(.*)\.\d{3}$', name):
        return name[:-4]
    return name


def is_node_groups_duplicate(node_groups):
    node_group_types = sorted([n.type for n in node_groups[0].nodes])
    return all( sorted([n.type for n in ng.nodes]) == 
             node_group_types for ng in node_groups[1:])


def is_images_duplicate(images):
    return all( img.filepath == images[0].filepath for image in images)


def is_materials_duplicate(materials):
    node_group_types = sorted([n.type for n in materials[0].node_tree.nodes])
    return all( sorted([n.type for n in mat.node_tree.nodes]) == 
             node_group_types for mat in materials[1:])


def merge_datablock_duplicates(datablocks, blend_data, force=False):
    """Merging materials, node_groups or images based on name .001, .002"""
    
    failed = []
    merged = []

    datablocks = list(datablocks)

    #blend_data = get_asset_type(datablocks[0].bl_rna.identifier)
    if blend_data == 'materials':
        is_datablock_duplicate = is_materials_duplicate
    elif blend_data == 'node_groups':
        is_datablock_duplicate = is_node_groups_duplicate
    elif blend_data == 'images':
        is_datablock_duplicate = is_images_duplicate
    else:
        raise Exception(f'Type, {blend_data} not supported')

    # Group by name
    groups = {}
    for datablock in images:
        groups.setdefault(clean_name(datablock.name), []).append(datablock)

    for datablock in blend_data:
        name = clean_name(datablock.name)
        if name in groups and datablock not in groups[name]:   
            groups[name].append(datablock)

    print("\nMerge Duplicate Datablocks...")

    for group in groups.values():
        if len(group) == 1:
            continue

        if not force:
            datablocks.sort(key=lambda x : x.name, reverse=True)

        for datablock in datablocks[1:]:
            is_duplicate = is_datablock_duplicate((datablock, datablocks[0]))
            
            if not is_duplicate and not force:
                failed.append((datablock.name, datablocks[0].name))
                print(f'Cannot merge Datablock {datablocks.name} with {datablocks[0].name} they are different')
                continue

            merged.append((datablock.name, datablocks[0].name))
            print(f'Merge Datablock {datablock.name} into {datablocks[0].name}')
            
            datablock.user_remap(datablocks[0])
            datablocks.remove(datablock)
            blend_data.remove(datablock)
    
    # Rename groups if it has no duplicate left
    for datablocks in groups.values():
        if len(datablocks) == 1 and not datablocks[0].library:
            datablocks[0].name = clean_name(datablocks[0].name)
    
    return merged, failed