From 17d494c1a2cb5229b764daeb415010f97fe46ad4 Mon Sep 17 00:00:00 2001
From: Roy Nieterau <roy_nieterau@hotmail.com>
Date: Tue, 5 Sep 2023 23:13:39 +0200
Subject: [PATCH] Add logic to collect user defined attributes and merge logic
 with pointcache and animation family + optimize the query by doing only one
 `cmds.listAttr` call

---
 .../maya/plugins/publish/collect_animation.py | 14 -------
 .../plugins/publish/collect_pointcache.py     | 15 -------
 .../collect_user_defined_attributes.py        | 39 +++++++++++++++++++
 3 files changed, 39 insertions(+), 29 deletions(-)
 create mode 100644 openpype/hosts/maya/plugins/publish/collect_user_defined_attributes.py

diff --git a/openpype/hosts/maya/plugins/publish/collect_animation.py b/openpype/hosts/maya/plugins/publish/collect_animation.py
index 8f523f770b..26a0a01c8b 100644
--- a/openpype/hosts/maya/plugins/publish/collect_animation.py
+++ b/openpype/hosts/maya/plugins/publish/collect_animation.py
@@ -58,17 +58,3 @@ class CollectAnimationOutputGeometry(pyblish.api.InstancePlugin):
         if instance.data.get("farm"):
             instance.data["families"].append("publish.farm")
 
-        # Collect user defined attributes.
-        if not instance.data.get("includeUserDefinedAttributes", False):
-            return
-
-        user_defined_attributes = set()
-        for node in hierarchy:
-            attrs = cmds.listAttr(node, userDefined=True) or list()
-            shapes = cmds.listRelatives(node, shapes=True) or list()
-            for shape in shapes:
-                attrs.extend(cmds.listAttr(shape, userDefined=True) or list())
-
-            user_defined_attributes.update(attrs)
-
-        instance.data["userDefinedAttributes"] = list(user_defined_attributes)
diff --git a/openpype/hosts/maya/plugins/publish/collect_pointcache.py b/openpype/hosts/maya/plugins/publish/collect_pointcache.py
index bb9065792f..5578a57f31 100644
--- a/openpype/hosts/maya/plugins/publish/collect_pointcache.py
+++ b/openpype/hosts/maya/plugins/publish/collect_pointcache.py
@@ -45,18 +45,3 @@ class CollectPointcache(pyblish.api.InstancePlugin):
         if proxy_set:
             instance.remove(proxy_set)
             instance.data["setMembers"].remove(proxy_set)
-
-        # Collect user defined attributes.
-        if not instance.data.get("includeUserDefinedAttributes", False):
-            return
-
-        user_defined_attributes = set()
-        for node in instance:
-            attrs = cmds.listAttr(node, userDefined=True) or list()
-            shapes = cmds.listRelatives(node, shapes=True) or list()
-            for shape in shapes:
-                attrs.extend(cmds.listAttr(shape, userDefined=True) or list())
-
-            user_defined_attributes.update(attrs)
-
-        instance.data["userDefinedAttributes"] = list(user_defined_attributes)
diff --git a/openpype/hosts/maya/plugins/publish/collect_user_defined_attributes.py b/openpype/hosts/maya/plugins/publish/collect_user_defined_attributes.py
new file mode 100644
index 0000000000..4d0790ad7c
--- /dev/null
+++ b/openpype/hosts/maya/plugins/publish/collect_user_defined_attributes.py
@@ -0,0 +1,39 @@
+from maya import cmds
+
+import pyblish.api
+
+
+class CollectUserDefinedAttributes(pyblish.api.InstancePlugin):
+    """Collect user defined attributes for nodes in instance."""
+
+    order = pyblish.api.CollectorOrder + 0.4
+    families = ["pointcache", "animation", "usd"]
+    label = "Collect User Defined Attributes"
+    hosts = ["maya"]
+
+    def process(self, instance):
+
+        # Collect user defined attributes.
+        if not instance.data.get("includeUserDefinedAttributes", False):
+            return
+
+        if "out_hierarchy" in instance.data:
+            # animation family
+            nodes = instance.data["out_hierarchy"]
+        else:
+            nodes = instance[:]
+        if not nodes:
+            return
+
+        shapes = cmds.listRelatives(nodes, shapes=True, fullPath=True) or []
+        nodes = set(nodes).union(shapes)
+
+        attrs = cmds.listAttr(list(nodes), userDefined=True) or []
+        user_defined_attributes = list(sorted(set(attrs)))
+        instance.data["userDefinedAttributes"] = user_defined_attributes
+
+        self.log.debug(
+            "Collected user defined attributes: {}".format(
+                ", ".join(user_defined_attributes)
+            )
+        )