wtd-libreflow-addon/scene_manager/library.py

"""Blender side of the Scene Manager: the assets linked from Libreflow publishes.

Every loaded publish is one ``bpy.types.Library``, tagged with custom
properties (``wtd_asset_oid``, ``wtd_file_oid``, ``wtd_revision``...). The tags
live on the library only: library overrides lose custom properties when they
resync. Instances are found back from their library: a collection instance
(empty) through its instanced collection, a library override through the
linked collection it overrides.
"""

import os
import re

import bpy

from .. import config

ASSET_OID = "wtd_asset_oid"


class LoadError(Exception):
    """An asset cannot be loaded or updated. The message is shown to the user."""


# --------------------------------------------------------------------------- #
# Loaded assets
# --------------------------------------------------------------------------- #
def managed_libraries():
    """Libraries loaded by the Scene Manager, sorted by asset name."""
    libraries = [lib for lib in bpy.data.libraries if lib.get(ASSET_OID)]
    return sorted(libraries, key=lambda lib: (lib.get("wtd_asset_name", ""), lib.name))


def find_library(file_oid):
    return next((lib for lib in managed_libraries() if lib.get("wtd_file_oid") == file_oid), None)


def linked_collection(library):
    """The asset collection linked from the library."""
    name = library.get("wtd_collection", "")
    return next((c for c in bpy.data.collections if c.library == library and c.name == name), None)


def all_instances():
    """{library name: [instances]}: collection instances (empties) and override hierarchy roots,
    found in one pass over the file."""
    found = {}
    for ob in bpy.data.objects:
        collection = ob.instance_collection if ob.instance_type == "COLLECTION" else None
        if ob.library is None and collection is not None and collection.library is not None:
            found.setdefault(collection.library.name, []).append(ob)
    for collection in bpy.data.collections:
        override = collection.override_library
        if (
            collection.library is None
            and override is not None
            and override.reference is not None
            and override.reference.library is not None
            and override.hierarchy_root == collection
        ):
            found.setdefault(override.reference.library.name, []).append(collection)
    return found


def instances(library):
    return all_instances().get(library.name, [])


def instance_objects(instance):
    return [instance] if isinstance(instance, bpy.types.Object) else list(instance.all_objects)


def selected_instances(view_layer):
    """(library, instance) pairs with at least one selected object."""
    selected = {ob.as_pointer() for ob in view_layer.objects if ob.select_get(view_layer=view_layer)}
    found = all_instances()
    return [
        (library, instance)
        for library in managed_libraries()
        for instance in found.get(library.name, [])
        if any(ob.as_pointer() in selected for ob in instance_objects(instance))
    ]


# --------------------------------------------------------------------------- #
# Load
# --------------------------------------------------------------------------- #
def link_mode(asset_type):
    modes = config.SCENE_MANAGER_LINK_MODES
    return modes.get(asset_type, modes.get("*", "INSTANCE"))


def _type_collection(scene, asset_type):
    """Collection gathering the instances of an asset type, under the scene collection."""
    name = config.SCENE_MANAGER_TYPE_COLLECTION.format(asset_type=asset_type)
    collection = bpy.data.collections.get(name)
    if collection is None or collection.library is not None:
        collection = bpy.data.collections.new(name)
    if collection.name not in scene.collection.children:
        scene.collection.children.link(collection)
    return collection


def _revision_path(revision, what):
    path = (revision or {}).get("path")
    if not path:
        raise LoadError(f"{what}: Libreflow gives no path for {revision and revision.get('name')}")
    if not os.path.isfile(path):
        raise LoadError(f"{what}: {revision['name']} not found on disk ({path})")
    return path


def link_asset(asset, revision, context_oid):
    """Link the asset collection of a published revision. Returns its library.

    The collection is the first of config SCENE_MANAGER_COLLECTION_NAMES found in
    the file, else its single top-level collection.
    """
    name = asset["name"]
    path = _revision_path(revision, name)
    candidates = [pattern.format(asset=name) for pattern in config.SCENE_MANAGER_COLLECTION_NAMES]
    with bpy.data.libraries.load(path, link=True) as (data_from, data_to):
        chosen = next((c for c in candidates if c in data_from.collections), None)
        data_to.collections = [chosen] if chosen else list(data_from.collections)
    linked = [c for c in data_to.collections if c is not None]
    if not linked:
        raise LoadError(f"{name}: no collection in {os.path.basename(path)}")
    library = linked[0].library

    if chosen is None:
        roots = [c for c in linked if not any(c in parent.children.values() for parent in linked)]
        if len(roots) != 1:
            message = (
                f"{name}: no collection named {' or '.join(candidates)}, and "
                f"{len(roots)} top-level collections in the file ({', '.join(c.name for c in roots)})"
            )
            bpy.data.libraries.remove(library)
            raise LoadError(message)
        chosen = roots[0].name

    library.name = name
    library[ASSET_OID] = asset["oid"]
    library["wtd_asset_name"] = name
    library["wtd_asset_type"] = asset.get("asset_type") or ""
    library["wtd_file_oid"] = asset["file"]["file_oid"]
    library["wtd_revision"] = revision["name"]
    library["wtd_collection"] = chosen
    library["wtd_context_oid"] = context_oid
    return library


def add_instance(scene, view_layer, library, parent=None):
    """Bring the library's asset in the scene: a library override or a collection instance
    (config SCENE_MANAGER_LINK_MODES), in ``parent`` or in the collection of its asset type."""
    collection = linked_collection(library)
    if collection is None:
        raise LoadError(f"{library.name}: collection {library.get('wtd_collection')} missing from the file")
    parent = parent or _type_collection(scene, library.get("wtd_asset_type", ""))
    if link_mode(library.get("wtd_asset_type", "")) == "OVERRIDE":
        override = collection.override_hierarchy_create(scene, view_layer)
        if scene.collection.children.get(override.name) == override and parent != scene.collection:
            scene.collection.children.unlink(override)
        if parent.children.get(override.name) != override:
            parent.children.link(override)
        return override
    empty = bpy.data.objects.new(library.get("wtd_asset_name", collection.name), None)
    empty.instance_type = "COLLECTION"
    empty.instance_collection = collection
    parent.objects.link(empty)
    return empty


# --------------------------------------------------------------------------- #
# Update / remove
# --------------------------------------------------------------------------- #
def update(library_name, revision):
    """Point a library to another revision of its file and reload it."""
    library = bpy.data.libraries[library_name]
    path = _revision_path(revision, library.get("wtd_asset_name", library_name))
    library.filepath = path
    library.reload()
    # The reload replaces the linked data: look the library up again.
    library = bpy.data.libraries[library_name]
    library["wtd_revision"] = revision["name"]
    return library


def remove_instance(instance):
    if isinstance(instance, bpy.types.Object):
        bpy.data.objects.remove(instance)
        return
    ids = {instance}
    ids.update(c for c in instance.children_recursive if c.override_library and c.library is None)
    for ob in instance.all_objects:
        if ob.override_library and ob.library is None:
            ids.add(ob)
            if ob.data is not None and ob.data.override_library and ob.data.library is None:
                ids.add(ob.data)
    bpy.data.batch_remove(ids)


def remove_library(library):
    """Remove every instance of the asset, then the library and its linked data."""
    for instance in instances(library):
        remove_instance(instance)
    bpy.data.libraries.remove(library)


# --------------------------------------------------------------------------- #
# Assetize
# --------------------------------------------------------------------------- #
def _holds_linked_data(collection):
    return any(c.override_library or c.library for c in collection.children_recursive) or any(
        ob.override_library or ob.library or (ob.instance_collection and ob.instance_collection.library)
        for ob in collection.all_objects
    )


def assetize_candidates():
    """Local collections that can be published: not linked, not an override, holding no loaded asset."""
    return [
        c
        for c in bpy.data.collections
        if c.library is None and c.override_library is None and not c.get(ASSET_OID) and not _holds_linked_data(c)
    ]


def asset_name_from(collection_name):
    """'COL_street lamp.001' -> 'street_lamp': the prefix and duplicate suffix removed, an identifier."""
    name = re.sub(r"\.\d{3,}$", "", collection_name)
    prefix = config.DATABLOCK_PREFIXES.get("collections", "")
    if prefix and name.upper().startswith(prefix.upper()):
        name = name[len(prefix):]
    return re.sub(r"[^A-Za-z0-9_]+", "_", name).strip("_")


def write_collection(collection, path):
    """Write the collection and what it uses to a .blend, for publishing."""
    bpy.data.libraries.write(path, {collection}, path_remap="RELATIVE_ALL", fake_user=True)


def parents_of(collection, scene):
    return [p for p in [scene.collection, *bpy.data.collections] if p.children.get(collection.name) == collection]