From a61f7ac7998453dad917d01e43fa134320a8e7e5 Mon Sep 17 00:00:00 2001
From: Roy Nieterau <roy_nieterau@hotmail.com>
Date: Wed, 6 Sep 2023 00:02:39 +0200
Subject: [PATCH] Add support for Look Assigner to assign looks in `aiStandin`
 for USD files based on `cbId` attributes in the USD file.

- For this to currently work the transform and shape should *not* be merged into a single Prim inside USD because otherwise the unique `cbId` between Transform and Shape node will be lost.
---
 .../tools/mayalookassigner/arnold_standin.py  |  8 ++++
 .../hosts/maya/tools/mayalookassigner/usd.py  | 38 +++++++++++++++++++
 2 files changed, 46 insertions(+)
 create mode 100644 openpype/hosts/maya/tools/mayalookassigner/usd.py

diff --git a/openpype/hosts/maya/tools/mayalookassigner/arnold_standin.py b/openpype/hosts/maya/tools/mayalookassigner/arnold_standin.py
index 0ce2b21dcd..076b0047bb 100644
--- a/openpype/hosts/maya/tools/mayalookassigner/arnold_standin.py
+++ b/openpype/hosts/maya/tools/mayalookassigner/arnold_standin.py
@@ -10,6 +10,7 @@ from openpype.client import get_last_version_by_subset_name
 from openpype.hosts.maya import api
 from . import lib
 from .alembic import get_alembic_ids_cache
+from .usd import is_usd_lib_supported, get_usd_ids_cache
 
 
 log = logging.getLogger(__name__)
@@ -74,6 +75,13 @@ def get_nodes_by_id(standin):
         # Support alembic files directly
         return get_alembic_ids_cache(path)
 
+    elif (
+        is_usd_lib_supported and
+        any(path.endswith(ext) for ext in [".usd", ".usda", ".usdc"])
+    ):
+        # Support usd files directly
+        return get_usd_ids_cache(path)
+
     json_path = None
     for f in os.listdir(os.path.dirname(path)):
         if f.endswith(".json"):
diff --git a/openpype/hosts/maya/tools/mayalookassigner/usd.py b/openpype/hosts/maya/tools/mayalookassigner/usd.py
new file mode 100644
index 0000000000..beecbd531a
--- /dev/null
+++ b/openpype/hosts/maya/tools/mayalookassigner/usd.py
@@ -0,0 +1,38 @@
+from collections import defaultdict
+
+try:
+    from pxr import Usd
+    is_usd_lib_supported = True
+except ImportError:
+    is_usd_lib_supported = False
+
+
+def get_usd_ids_cache(path):
+    # type: (str) -> dict
+    """Build a id to node mapping in a USD file.
+
+    Nodes without IDs are ignored.
+
+    Returns:
+        dict: Mapping of id to nodes in the USD file.
+
+    """
+    if not is_usd_lib_supported:
+        raise RuntimeError("No pxr.Usd python library available.")
+
+    stage = Usd.Stage.Open(path)
+    ids = {}
+    for prim in stage.Traverse():
+        attr = prim.GetAttribute("userProperties:cbId")
+        if not attr.IsValid():
+            continue
+        path = str(prim.GetPath())
+        value = attr.Get()
+        if not value:
+            continue
+        ids[path] = value
+
+    cache = defaultdict(list)
+    for path, value in ids.items():
+        cache[value].append(path)
+    return dict(cache)