asset_library/file/bundle.py

import argparse
import sys
import json
from pathlib import Path
import bpy
import re
import uuid
from itertools import groupby

from asset_library.constants import ASSETLIB_FILENAME, MODULE_DIR
from asset_library.common.bl_utils import thumbnail_blend_file
from asset_library.common.functions import command


@command
def bundle_library(
    source_directory,
    bundle_directory,
    template_info,
    thumbnail_template,
    template=None,
    data_file=None,
):

    field_pattern = r"{(\w+)}"
    asset_data_path = Path(bundle_directory) / ASSETLIB_FILENAME

    glob_pattern = re.sub(field_pattern, "*", template)
    re_pattern = re.sub(field_pattern, r"([\\w -_.]+)", template)
    re_pattern = re_pattern.replace("?", ".")

    field_names = re.findall(field_pattern, template)

    asset_file_datas = []
    for f in sorted(Path(source_directory).glob(glob_pattern)):
        rel_path = f.relative_to(source_directory).as_posix()

        field_values = re.findall(re_pattern, rel_path)[0]
        field_data = {k: v for k, v in zip(field_names, field_values)}

        name = field_data.get("name", f.stem)
        thumbnail = (f / thumbnail_template.format(name=name)).resolve()
        asset_data = (f / template_info.format(name=name)).resolve()

        catalogs = sorted(
            [v for k, v in sorted(field_data.items()) if re.findall("cat[0-9]+", k)]
        )
        catalogs = [c.replace("_", " ").title() for c in catalogs]

        if not thumbnail.exists():
            thumbnail_blend_file(f, thumbnail)

        asset_data = {
            "catalog": "/".join(catalogs),
            "preview": thumbnail.as_posix(),  #'./' + bpy.path.relpath(str(thumbnail), start=str(f))[2:],
            "filepath": f.as_posix(),  #'./' + bpy.path.relpath(str(f), start=str(asset_data_path))[2:],
            "name": name,
            "tags": [],
            "metadata": {"filepath": f.as_posix()},
        }

        asset_file_datas.append(asset_data)

    # Write json data file to store all asset found
    print(f"Writing asset data file to, {asset_data_path}")
    asset_data_path.write_text(json.dumps(asset_file_datas, indent=4))

    # script = MODULE_DIR / 'common' / 'bundle_blend.py'
    # cmd = [bpy.app.binary_path, '--python', str(script), '--', '--filepath', str(filepath)]
    # print(cmd)
    # subprocess.call(cmd)


@command
def bundle_blend(filepath, depth=0):
    # print('Bundle Blend...')
    filepath = Path(filepath)

    # asset_data_path = get_asset_datas_file(filepath)

    asset_data_path = filepath / ASSETLIB_FILENAME
    blend_name = filepath.name.replace(" ", "_").lower()
    blend_path = (filepath / blend_name).with_suffix(".blend")

    if not asset_data_path.exists():
        raise Exception(f"The file {asset_data_path} not exist")

    catalog_path = get_catalog_path(filepath)
    catalog_data = read_catalog(catalog_path)

    asset_file_data = json.loads(asset_data_path.read_text())
    # asset_file_data = {i['catalog']:i for i in asset_file_data}

    if depth == 0:
        groups = [asset_file_data]
    else:
        asset_file_data.sort(key=lambda x: x["catalog"].split("/")[:depth])
        groups = groupby(asset_file_data, key=lambda x: x["catalog"].split("/")[:depth])

    # progress = 0
    total_assets = len(asset_file_data)

    i = 0
    for sub_path, asset_datas in groups:
        bpy.ops.wm.read_homefile(use_empty=True)

        for asset_data in asset_datas:
            blend_name = sub_path[-1].replace(" ", "_").lower()
            blend_path = Path(filepath, *sub_path, blend_name).with_suffix(".blend")

            if i % int(total_assets / 100) == 0:
                print(f"Progress: {int(i / total_assets * 100)}")

            col = bpy.data.collections.new(name=asset_data["name"])

            # Seems slow
            # bpy.context.scene.collection.children.link(col)
            col.asset_mark()

            with bpy.context.temp_override(id=col):
                bpy.ops.ed.lib_id_load_custom_preview(filepath=asset_data["preview"])

            col.asset_data.description = asset_data.get("description", "")

            catalog_name = asset_data["catalog"]
            catalog = catalog_data.get(catalog_name)
            if not catalog:
                catalog = {"id": str(uuid.uuid4()), "name": catalog_name}
                catalog_data[catalog_name] = catalog

            col.asset_data.catalog_id = catalog["id"]

            for k, v in asset_data.get("metadata", {}).items():
                col.asset_data[k] = v

            i += 1

        print(f"Saving Blend to {blend_path}")

        blend_path.mkdir(exist_ok=True, parents=True)
        bpy.ops.wm.save_as_mainfile(filepath=str(blend_path), compress=True)

    write_catalog(catalog_path, catalog_data)

    bpy.ops.wm.quit_blender()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="bundle_blend",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )

    parser.add_argument("--source-path")
    parser.add_argument("--bundle-path")
    parser.add_argument("--asset-data-template")
    parser.add_argument("--thumbnail-template")
    parser.add_argument("--template", default=None)
    parser.add_argument("--data-file", default=None)
    parser.add_argument("--depth", default=0, type=int)

    if "--" in sys.argv:
        index = sys.argv.index("--")
        sys.argv = [sys.argv[index - 1], *sys.argv[index + 1 :]]

    args = parser.parse_args()

    bundle_library(
        source_directory=args.source_directory,
        bundle_directory=args.bundle_directory,
        template_info=args.template_info,
        thumbnail_template=args.thumbnail_template,
        template=args.template,
        data_file=args.data_file,
    )

    bundle_blend(filepath=args.bundle_directory, depth=args.depth)