Maya Shader Exporter

export_shaders.py

"""
This is an attempt to make a poor man's version of Katana's look files.

We record any shader assignments on the descendants of a given group, each location is recorded as a relative path along
with the shading engine that is assigned to it.

Encoding:

    given the following nodes:

    root|bob|geo|head|eye_l
    root|bob|geo|head|eye_r

    using root|character as the root group, we collect the following assignments:

    [
        ("|geo|head|eye_l", "eyes_SG"),
        ("|geo|head|eye_r", "eyes_SG"),
    ]

    These are saved as json encoded string attributes on a node which is by default called __assignments__ and saved
    along with the required shading engines, and any connected shading nodes.

Referencing:
    we reference this shader file into our scene under the looks namespace, giving it its own sub namespace based on
    asset name.

    e.g. looks:bob

    under this namespace we have

    looks:bob:__assignments__
    looks:bob:eyes_SG
    looks_bob:eyes_shader
    looks_bob:eyes_texture

    this keeps each look self-contained and easy to manage.

Application:
    we import an animated instance of bob into our maya scene, and put it under a named group

    e.g.

    |root|chars|bob_01

    under which we have

    |root|chars|bob_01|geo|head|eye_l
    |root|chars|bob_01|geo|head|eye_r
    ...

    to apply our bob 'look' we do the following steps for each assignment

    - prefix the relative path with our root group so |geo|head_eye_1 becomes: |root|chars|bob_01|geo|head|eye_l
    - prefix the shading group name with our look namespace so eyes_SG becomes:  looks:bob:eyes_SG
    - assign the shading group looks:bob:eyes_SG to |root|chars|bob_01|geo|head|eye_l

In combination with alembic files this allows externalising of both geometry and shaders to separate files, and makes
updating them trivial providing the hierarchies remain fixed.

"""

import json
import os
import typing
import logging

from maya import cmds

logger = logging.getLogger(__name__)

DEFAULT_LOOKS_NAMESPACE = "looks"
DEFAULT_ASSIGNMENTS_GROUP = "__assignments__"


def collect_assignments(root_group):
    """
    Collect the shader assignments from the given root group.

    This method walks the hierarchy of the root group and returns a list of the shader assignments for each transform
    with a shape / shader assigned to it.

    Assignments are returned as a list of tuples of the form:

    [
        ("|group1|group2|pSphere1", "initialShadingGroup"),

        ("|group1|group2|pSphere2", "initialShadingGroup"),

    ]


    Note:
        This function assumes that the root group is a group that contains all the geometry for a single asset, and
        that shader assignments are not face-sets.

    Args:
        root_group(str): The root group to get the shader assignments for.

    Returns:
        (typing.List[typing.Tuple[str,str]]): A list of tuples of the form (relative_name, shading_engine)

    """
    long_root_group = cmds.ls(root_group, long=True)[0]

    assignments = []

    for descendant in cmds.listRelatives(long_root_group, allDescendents=True, fullPath=True):
        if cmds.nodeType(descendant) == "shape":
            continue

        try:
            shape = cmds.listRelatives(descendant, shapes=True, fullPath=True)[0]
        except TypeError:
            continue

        try:
            shading_engine = cmds.listConnections(shape, type="shadingEngine")[0]
        except TypeError:
            continue

        rel_name = descendant[len(long_root_group):]
        assignments.append((rel_name, shading_engine))

    return assignments


def encode_assignments(assignments_group, assignments):
    """ Encode the given assignments into the given assignments group.

    An assignment is a tuple of the relative name of the node and the shading engine it is assigned to.
    e.g. ("|group1|group2|pSphere1", "initialShadingGroup")

    Each assignment is encoded as a json string and stored as an attribute on the assignments group. A new attribute
    is created for each assignment to avoid the 16k limit on string attributes when encoding large scenes.

    Args:
        assignments_group(str): The group to encode the assignments on.
        assignments(typing.Iterable[typing.Tuple[str, str]]): The assignments to encode.

    """
    for index, (rel_name, shading_engine) in enumerate(assignments):
        cmds.addAttr(assignments_group, longName=f"assignment_{index}", dataType="string")
        cmds.setAttr(f"{assignments_group}.assignment_{index}", json.dumps((rel_name, shading_engine)), type="string")


def decode_assignments(assignment_group):
    """
    Get the assignments from the given assignment group.

    Note:
        this method assumes that attribute indexes are sequential and start at 0. If this is not the case, then this
        method will not return all the assignments.

    Args:
        assignment_group(str): The assignment group to get the assignments from.

    Returns:
        typing.Dict[str, str]: A dictionary mapping the relative name of the node to the shading engine it is assigned
            to.
    """
    assignments = {}

    index = 0
    while True:
        try:
            assignment_string = cmds.getAttr(f"{assignment_group}.assignment_{index}")
        except ValueError:
            break

        rel_name, shading_engine = json.loads(assignment_string)
        assignments[rel_name] = shading_engine
        index += 1

    return assignments


def export_shading_engines(shading_engines, output_file, include_nodes=None):
    """
    Export the given shading engines to the given file.

    This method uses a quirk of export selected to export the shading engines to the given file. If you export a mesh
    with a shading engine assigned to it, the shading engine will be exported as well. However, if you export a shading
    engine, it will export all the meshes that are assigned to it, and everything else connected to those meshes.

    This method takes advantage of this quirk by creating a cube for each unique shading engine and exporting a minimal
    maya file. While not ideal this is better than the alternative of exporting the entire scene and then pruning out
    the unneeded nodes.

    Args:
        shading_engines(typing.Iterable[str]): The shading engines to export.
        output_file(str): The file to export the shading engines to.
        include_nodes(typing.Iterable[str]): Additional nodes to include in the export.
    """
    shading_engines = set(shading_engines)

    anchors_group = cmds.group(empty=True, name="__anchors__")
    anchors = []
    for shading_engine in shading_engines:
        cube = cmds.polyCube()[0]
        cmds.sets(cube, edit=True, forceElement=shading_engine)
        cmds.parent(cube, anchors_group)
        cmds.setAttr(f"{cube}.visibility", False)
        anchors.append(cube)

    # export the file
    ext = output_file.rsplit(".")[-1]
    if ext == "ma":
        filetype = "mayaAscii"
    elif ext == "mb":
        filetype = "mayaBinary"
    else:
        raise ValueError(f"Unsupported extension: {ext}")

    selection = [*anchors]

    if isinstance(include_nodes, (list, tuple, set)):
        selection.extend(include_nodes)

    cmds.select(selection, replace=True)
    cmds.file(output_file, exportSelected=True, force=True, type=filetype)
    cmds.delete(anchors_group)


def export_look_file(filepath, root_group, assignments_group_name=None):
    """
    Export the look file for the given root group.

    Args:
        filepath(str): The path to export the look file to.
        root_group(str): The root group to export the look file for.
        assignments_group_name(str): The name of the group to store the assignments in.
    """
    assignments_group_name = assignments_group_name or DEFAULT_ASSIGNMENTS_GROUP

    long_root_group = cmds.ls(root_group, long=True)[0]

    # create a transform called __assignments__
    if cmds.objExists(assignments_group_name):
        cmds.delete(assignments_group_name)

    assignment_group = cmds.createNode("transform", name=assignments_group_name)

    # collect the assignments
    assignments = collect_assignments(long_root_group)

    # encode the assignments
    encode_assignments(assignment_group, assignments)

    # export the shading engines
    shading_engines = set([shading_engine for _, shading_engine in assignments])
    export_shading_engines(shading_engines, filepath, include_nodes=[assignment_group])
    cmds.delete(assignment_group)


def import_look_file(filepath, looks_namespace=None, look_name=None, delete_existing=True):
    """
    Import the look file at the given path and store it in the namespace: looks:look_name
    """
    looks_namespace = looks_namespace or DEFAULT_LOOKS_NAMESPACE
    look_name = look_name or os.path.basename(filepath).rsplit(".", 1)[0]

    if not cmds.namespace(exists=looks_namespace):
        cmds.namespace(addNamespace=looks_namespace)

    import_namespace = f"{looks_namespace}:{look_name}"

    if cmds.namespace(exists=import_namespace):
        if delete_existing:
            cmds.namespace(removeNamespace=import_namespace, deleteNamespaceContent=True)
        else:
            raise ValueError(f"look already exists: {import_namespace}")

    cmds.file(filepath, reference=True, namespace=import_namespace)


def apply_look(root_group, look, looks_namespace=None, error_on_missing=True, assignments_group_name=None):
    looks_namespace = looks_namespace or DEFAULT_LOOKS_NAMESPACE
    assignments_group_name = assignments_group_name or DEFAULT_ASSIGNMENTS_GROUP

    namespace = f"{looks_namespace}:{look}"
    assignment_group = f"{namespace}:{assignments_group_name}"
    assignments = decode_assignments(assignment_group)

    long_root_group = cmds.ls(root_group, long=True)[0]
    for rel_name, shading_engine in assignments.items():
        abs_node_name = f"{long_root_group}{rel_name}"
        abs_shading_engine_name = f"{namespace}:{shading_engine}"

        if not all([cmds.objExists(abs_node_name), cmds.objExists(abs_shading_engine_name)]):
            if error_on_missing:
                raise ValueError(f"Missing node or shading engine: {abs_node_name} {abs_shading_engine_name}")
            else:
                logger.warning(f"Missing node or shading engine: {abs_node_name} {abs_shading_engine_name}")

        cmds.sets(abs_node_name, edit=True, forceElement=abs_shading_engine_name)


def get_looks(looks_namespace=None):
    looks_namespace = looks_namespace or DEFAULT_LOOKS_NAMESPACE

    if not cmds.namespace(exists=looks_namespace):
        return []

    return cmds.namespaceInfo(looks_namespace, listOnlyNamespaces=True)


filepath = "~/Desktop/shaders.mb"
# export_look_file(filepath, "geo_all")
import_look_file(filepath, look_name="bee")
# apply_look("geo_all", "test")