Maya Shader Exporter

export_shaders.py

"""
Rob Fletcher 2025

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 os
from PySide6 import QtCore, QtWidgets, QtGui

import fnmatch
import json
import logging

from maya import cmds

logger = logging.getLogger(__name__)

__HANDLE__ = None

DEFAULT_LOOKS_NAMESPACE = "looks"
DEFAULT_ASSIGNMENTS_GROUP = "__assignments__"
# node type, attribute name mask.
ATTRIBUTE_FILTERS = [
    (
        "shape",
        "ai*",
    ),
    ("shape", "visibility"),
    ("shape", "primaryVisibility"),
    ("shape", "castsShadows"),
]

# attributes which don't define their default correctly.
UNDEFINED_DEFAULTS = {
    "aiMotionVectorSource": "velocityPV",
    "aiAov": "default",
    "aiTranslator": "polymesh",
}


def abs_path(root_group, relpath):
    """
    builds an absolute path from the given root group and relative path. takes into account namespaces.

    Examples:

        abs_path("root:group", "child") -> "root:group|root:group:child"
        abs_path("root:group", "child|grandchild") -> "root:group|root:group:child|root:group:child:grandchild"
        abs_path("root", "child") -> "root|child"

    Args:
        root_group(str): the root group to use as the base.
        relpath(str): the relative path to convert.

    Returns:
        str: the absolute path.

    """
    if ":" in root_group:
        namespace = root_group[: root_group.rindex(":")]
        rel_with_namespace = "|".join(f"{namespace}:{x}" for x in relpath.split("|"))
        return f"{root_group}|{rel_with_namespace}"
    else:
        return f"{root_group}|{relpath.lstrip('|')}"


def is_default_value(node, attr_name):
    """
    Checks if the given attribute on the given node is set to its default value.

    Args:
        node(str): the node to check.
        attr_name(str): the attribute to check.

    Returns:
        bool: True if the attribute is set to its default value, False otherwise.
    """

    # some attributes don't define their default value correctly, so we have to hard code them.
    # why is arnold so bad at this?
    if attr_name in UNDEFINED_DEFAULTS:
        default_value = UNDEFINED_DEFAULTS[attr_name]
    else:
        default_value = cmds.attributeQuery(attr_name, node=node, listDefault=True)

    attr_type = cmds.getAttr(f"{node}.{attr_name}", type=True)

    current_value = cmds.getAttr(f"{node}.{attr_name}")

    # default values are always returned as a list, so we need to unpack them, and convert them to the correct type.
    if attr_type == "bool":
        default_value = bool(default_value[0])
    elif attr_type in ["long", "short", "byte", "char", "enum"]:
        default_value = int(default_value[0])
    elif attr_type in ["float", "double"]:
        default_value = float(default_value[0])
    elif fnmatch.fnmatch(attr_type, "float*") or fnmatch.fnmatch(attr_type, "double*"):
        current_value = current_value[0]
    elif fnmatch.fnmatch(attr_type, "color*"):
        current_value = current_value[0]
    elif attr_type == "string" and default_value is None:
        default_value = ""

    # if we're dealing with a compound attribute then compare each of its values.
    if isinstance(current_value, (list, tuple)):
        return all(a == b for a, b in zip(current_value, default_value))  # type: ignore

    # otherwise do a simple comparison
    return current_value == default_value


def collect_attribute_assignments(root_group, attribute_filters=None):
    """
    Collects a list of non-default attribute assignments for the given root group. Only attributes which match the
    given attribute filters are collected. If no attribute filters are given then the default filters are used.

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

    [
        ("|group1|group2|pSphere1", "attribute_name", "attribute_value"),
        ("|group1|group2|pSphere1", "attribute_name", "attribute_value"),
    ]

    Args:
        root_group(str): the root group to collect assignments from.
        attribute_filters(list[tuple[str, str]]): a list of node type, attribute name patterns to collect assignments
            for. see ATTRIBUTE_FILTERS for an example.
    Returns:
        list[tuple[str, str, typing.Union[int,str,bool,float]]]: a list of (relpath, attribute, value) pairs.

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

    nodes = {}

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

        nodes.setdefault(cmds.nodeType(descendant), []).append(descendant)

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

        nodes.setdefault("shape", []).append(shape)

    attribute_assignments = []
    for node_type, attr_pattern in attribute_filters:
        for node in nodes.get(node_type, []):
            rel_name = node[len(long_root_group) :]

            for attr_name in cmds.listAttr(node):
                # skip attributes not matching the pattern
                if not fnmatch.fnmatch(attr_name, attr_pattern):
                    continue

                node_type = cmds.getAttr(f"{node}.{attr_name}", type=True)
                if node_type == "message":
                    continue

                if is_default_value(node, attr_name):
                    continue

                value = cmds.getAttr(f"{node}.{attr_name}")

                attribute_assignments.append((rel_name, attr_name, value))
    return attribute_assignments


def collect_shader_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, attribute_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, data in enumerate(assignments):
        cmds.addAttr(
            assignments_group, longName=f"assignment_{index}", dataType="string"
        )
        cmds.setAttr(
            f"{assignments_group}.assignment_{index}",
            json.dumps(data),
            type="string",
        )

    for index, data in enumerate(attribute_assignments):
        cmds.addAttr(
            assignments_group, longName=f"attribute_{index}", dataType="string"
        )

        cmds.setAttr(
            f"{assignments_group}.attribute_{index}",
            json.dumps(data),
            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.
    """
    shader_assignments = {}
    attribute_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)
        shader_assignments[rel_name] = shading_engine
        index += 1

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

        attribute_assignments.append(json.loads(assignment_string))
        index += 1

    return shader_assignments, attribute_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.
    """
    logger.info(f"exporting shading engines to: {output_file}")
    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)
    try:
        # 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)
    finally:
        cmds.delete(anchors_group)


def export_look(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

    # always use the full name of nodes this ensures all rel names are accurate.
    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
    shader_assignments = collect_shader_assignments(long_root_group)
    attribute_assignments = collect_attribute_assignments(root_group)

    # encode the assignments
    encode_assignments(assignment_group, shader_assignments, attribute_assignments)

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


def reference_look(look_name, filepath, looks_namespace=None):
    """
    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

    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):
        raise ValueError(f"namespace: {import_namespace} already exists.")

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


def replace_look(look_name, filepath, looks_namespace=None):
    looks_namespace = looks_namespace or DEFAULT_LOOKS_NAMESPACE
    look_namespace = f"{looks_namespace}:{look_name}"

    if not cmds.namespace(exists=looks_namespace):
        raise ValueError(f"namespace: {look_namespace} doesn't exist.")

    contents = cmds.namespaceInfo(looks_namespace, listOnlyNamespaces=True)
    if not contents:
        raise ValueError(f"namespace: {look_namespace} is empty.")

    reference = cmds.referenceQuery(contents[0], referenceNode=True)
    cmds.file(filepath, loadReference=reference)


def remove_look(look_name, looks_namespace=None):
    looks_namespace = looks_namespace or DEFAULT_LOOKS_NAMESPACE
    look_namespace = f"{looks_namespace}:{look_name}"

    if not cmds.namespace(exists=look_namespace):
        raise ValueError(f"namespace: {looks_namespace} doesn't exist")

    contents = cmds.namespaceInfo(looks_namespace, listOnlyNamespaces=True)
    if not contents:
        raise ValueError(f"namespace: {look_namespace} is empty.")

    filename = cmds.referenceQuery(contents[0], filename=True)
    cmds.file(filename, reference=True, removeReference=True)


def set_attribute_value(node, attr_name, value):
    if not cmds.attributeQuery(attr_name, node=node, exists=True):
        if isinstance(value, str):
            attr_type = "string"
        elif isinstance(value, float):
            attr_type = "float"
        elif isinstance(value, int):
            attr_type = "long"
        elif isinstance(value, bool):
            attr_type = "bool"
        else:
            raise ValueError(f"unsupported attribute type: {type(value)}")

        # Create the attribute
        cmds.addAttr(node, longName=attr_name, attributeType=attr_type)

    # Set the attribute value
    if isinstance(value, str):
        cmds.setAttr(node + "." + attr_name, value, type="string")
    else:
        cmds.setAttr(node + "." + attr_name, value)


def apply_look(
    root_group,
    look,
    looks_namespace=None,
    error_on_missing=True,
    assignments_group_name=None,
):
    """
    Apply the given look to the given root group.

    Args:
        root_group(str):

    """
    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}"

    shader_assignments, attribute_assignments = decode_assignments(assignment_group)

    long_root_group = cmds.ls(root_group, long=True)[0]
    for rel_name, shading_engine in shader_assignments.items():
        abs_node_name = abs_path(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)

    for rel_name, attribute_name, value in attribute_assignments:
        full_node_name = abs_path(long_root_group, rel_name)
        set_attribute_value(full_node_name, attribute_name, value)


def get_looks(looks_namespace=None):
    """
    Get a list of all the looks in the given namespace.
    """
    looks_namespace = looks_namespace or DEFAULT_LOOKS_NAMESPACE

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

    contents = cmds.namespaceInfo(looks_namespace, listOnlyNamespaces=True)

    if not contents:
        return []

    return [n.rsplit(":")[-1] for n in contents]


def export_look_action():
    """
    Export the look file for the currently selected root group.
    """
    selection = cmds.ls(selection=True) or []

    if len(selection) > 1:
        cmds.confirmDialog(message="Only one transform should be selected")
        return
    elif len(selection) == 0:
        cmds.confirmDialog(message="Please select the root geometry group of an asset.")
        return

    root_group = selection[0]

    filter_to_extension = {
        "Maya Binary (*.mb)": "mb",
        "Maya Ascii (*.ma)": "ma",
    }

    filename, filter_ = QtWidgets.QFileDialog.getSaveFileName(
        None, caption="Save Look File", filter=";;".join(filter_to_extension.keys())
    )

    name, _ = os.path.splitext(filename)
    filename = f"{name}.{filter_to_extension[filter_]}"

    if not filename:
        return

    name, ext = os.path.splitext(filename)
    if not ext:
        filename = f"{filename}.{filter_}"

    export_look(filename, root_group)


def load_look_action():
    """
    Load a look file into the current scene. If the look name already exists, it will replace the existing look.
    """
    filename, _ = QtWidgets.QFileDialog.getOpenFileName(
        None, caption="Load Look File", filter="Maya Files (*.ma *.mb)"
    )

    if not filename:
        return

    text = os.path.basename(filename).rsplit(".")[0]

    result = cmds.promptDialog(
        title="Look Name", message="Look Name:", button=["OK", "Cancel"], text=text
    )

    if result == "Cancel":
        return

    look_name = cmds.promptDialog(query=True, text=True)

    if look_name in get_looks():
        replace_look(look_name, filename)
    else:
        reference_look(look_name, filename)


class LookToolView(QtWidgets.QWidget):
    refreshClicked = QtCore.Signal()
    applyClicked = QtCore.Signal(QtCore.QModelIndex)
    loadClicked = QtCore.Signal()
    saveClicked = QtCore.Signal()
    removeClicked = QtCore.Signal()
    lookClicked = QtCore.Signal(QtCore.QModelIndex)

    opened = QtCore.Signal()
    closed = QtCore.Signal()

    def __init__(self, parent=None):
        super().__init__(parent=parent)
        self.setWindowTitle("Look Tool")
        self.look_list = QtWidgets.QListView()
        self.look_list.clicked.connect(self.onListViewClicked)

        self.apply_button = QtWidgets.QPushButton("Apply")
        self.apply_button.clicked.connect(self.onApplyClicked)
        self.apply_button.setEnabled(False)

        self.refresh_button = QtWidgets.QPushButton("Refresh")
        self.refresh_button.clicked.connect(self.refreshClicked)

        self.load_button = QtWidgets.QPushButton("Load")
        self.load_button.clicked.connect(self.loadClicked)

        self.save_button = QtWidgets.QPushButton("Save")
        self.save_button.clicked.connect(self.saveClicked)

        self.remove_button = QtWidgets.QPushButton("Remove")
        self.remove_button.clicked.connect(self.removeClicked)

        button_layout = QtWidgets.QHBoxLayout()
        button_layout.addWidget(self.save_button)
        button_layout.addWidget(self.load_button)
        button_layout.addWidget(self.remove_button)
        button_layout.addWidget(self.refresh_button)
        button_layout.addWidget(self.apply_button)

        layout = QtWidgets.QVBoxLayout(self)
        layout.addWidget(self.look_list)
        layout.addLayout(button_layout)

    def setModel(self, model):
        self.look_list.setModel(model)

    def onListViewClicked(self, index: QtCore.QModelIndex):
        self.lookClicked.emit(index)

    def onApplyClicked(self):
        index = self.look_list.currentIndex()
        self.applyClicked.emit(index)

    def setApplyEnabled(self, enabled):
        self.apply_button.setEnabled(enabled)

    def setSaveEnabled(self, enabled):
        self.save_button.setEnabled(enabled)

    def setRemoveEnabled(self, enabled):
        self.remove_button.setEnabled(enabled)

    def closeEvent(self, event: QtGui.QCloseEvent) -> None:
        super().closeEvent(event)
        self.closed.emit()

    def showEvent(self, event: QtGui.QShowEvent) -> None:
        super().showEvent(event)
        self.opened.emit()

    def selectedLook(self):
        return self.look_list.currentIndex().data(
            role=QtCore.Qt.ItemDataRole.DisplayRole
        )


class LookModel(QtCore.QAbstractListModel):
    def __init__(self, look_namespace=None, parent=None):
        super().__init__(parent=parent)
        self._items = []
        self._look_namespace = look_namespace

    def rowCount(self, parent: QtCore.QModelIndex = ...) -> int:
        return len(self._items)

    def data(
        self, index: QtCore.QModelIndex, role: int = QtCore.Qt.ItemDataRole.DisplayRole
    ):
        if not index.isValid():
            return

        if role == QtCore.Qt.ItemDataRole.DisplayRole:
            return self._items[index.row()]

    def refresh(self):
        self.beginResetModel()
        self._items = get_looks(looks_namespace=self._look_namespace)
        self.endResetModel()


class LookToolController(QtCore.QObject):
    def __init__(self, parent=None, looks_namespace=None):
        super().__init__(parent=parent)
        self._selection_changed_script_job = None
        self._new_scene_opened_script_job = None
        self._looks_namespace = looks_namespace

        self.model = LookModel(look_namespace=looks_namespace)

        self.view = LookToolView()
        self.view.setModel(self.model)

        self.view.refreshClicked.connect(self.model.refresh)
        self.view.saveClicked.connect(self.onSaveClicked)
        self.view.loadClicked.connect(self.onLoadClicked)
        self.view.applyClicked.connect(self.onApplyClicked)
        self.view.removeClicked.connect(self.onRemoveClicked)
        self.view.opened.connect(self.onViewOpened)
        self.view.closed.connect(self.onViewClosed)

        self.updateViewState()

    def updateViewState(self):
        nodes_selected = bool(cmds.ls(selection=True))
        look_selected = bool(self.view.selectedLook())

        self.view.setSaveEnabled(nodes_selected)
        self.view.setApplyEnabled(nodes_selected and look_selected)
        self.view.setRemoveEnabled(look_selected)

    def onViewOpened(self):
        self._selection_changed_script_job = cmds.scriptJob(
            event=["SelectionChanged", self.updateViewState]
        )
        self._new_scene_opened_script_job = cmds.scriptJob(
            event=["NewSceneOpened", self.model.refresh]
        )

    def onViewClosed(self):
        cmds.scriptJob(kill=self._selection_changed_script_job)
        cmds.scriptJob(kill=self._new_scene_opened_script_job)

    def onApplyClicked(self, index: QtCore.QModelIndex):
        look_name = index.data(role=QtCore.Qt.ItemDataRole.DisplayRole)

        for root_grp in cmds.ls(selection=True):
            apply_look(root_grp, look_name, looks_namespace=self._looks_namespace)

    def onRemoveClicked(self):
        look = self.view.selectedLook()
        remove_look(look, looks_namespace=self._looks_namespace)

    def onSaveClicked(self):
        export_look_action()

    def onLoadClicked(self):
        load_look_action()
        self.model.refresh()

    def execute(self):
        self.model.refresh()
        self.view.show()


LookToolController().execute()