From 7d7fd313fa7d75d6eb92f2cd1120b192b1ca8b41 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20Samohel?= <ondrej.samohel@ynput.io>
Date: Mon, 31 Mar 2025 15:42:06 +0200
Subject: [PATCH] :recycle: wrap trait based instance into functions

---
 client/ayon_core/pipeline/publish/__init__.py | 10 +++
 client/ayon_core/pipeline/publish/lib.py      | 71 ++++++++++++++++++-
 .../plugins/publish/integrate_traits.py       | 14 ++--
 3 files changed, 90 insertions(+), 5 deletions(-)

diff --git a/client/ayon_core/pipeline/publish/__init__.py b/client/ayon_core/pipeline/publish/__init__.py
index 5363e0b378..ede7fc3a35 100644
--- a/client/ayon_core/pipeline/publish/__init__.py
+++ b/client/ayon_core/pipeline/publish/__init__.py
@@ -46,6 +46,11 @@ from .lib import (
     get_publish_instance_families,
 
     main_cli_publish,
+
+    add_trait_representations,
+    get_trait_representations,
+    has_trait_representations,
+    set_trait_representations,
 )
 
 from .abstract_expected_files import ExpectedFiles
@@ -104,4 +109,9 @@ __all__ = (
 
     "RenderInstance",
     "AbstractCollectRender",
+
+    "add_trait_representations",
+    "get_trait_representations",
+    "has_trait_representations",
+    "set_trait_representations",
 )
diff --git a/client/ayon_core/pipeline/publish/lib.py b/client/ayon_core/pipeline/publish/lib.py
index 49ecab2221..f1dda288a6 100644
--- a/client/ayon_core/pipeline/publish/lib.py
+++ b/client/ayon_core/pipeline/publish/lib.py
@@ -6,7 +6,7 @@ import inspect
 import copy
 import warnings
 import xml.etree.ElementTree
-from typing import Optional, Union, List
+from typing import TYPE_CHECKING, Optional, Union, List
 
 import ayon_api
 import pyblish.util
@@ -27,6 +27,11 @@ from .constants import (
     DEFAULT_HERO_PUBLISH_TEMPLATE,
 )
 
+if TYPE_CHECKING:
+    from ayon_core.pipeline.traits import Representation
+
+
+TRAIT_INSTANCE_KEY: str = "representations_with_traits"
 
 def get_template_name_profiles(
     project_name, project_settings=None, logger=None
@@ -1062,3 +1067,67 @@ def main_cli_publish(
             sys.exit(1)
 
     log.info("Publish finished.")
+
+
+def has_trait_representations(
+        instance: pyblish.api.Instance) -> bool:
+    """Check if instance has trait representation.
+
+    Args:
+        instance (pyblish.api.Instance): Instance to check.
+    
+    Returns:
+        True: Instance has trait representation.
+        False: Instance does not have trait representation.
+    
+    """
+    return bool(instance.data.get(TRAIT_INSTANCE_KEY))
+
+
+def add_trait_representations(
+        instance: pyblish.api.Instance,
+        representations: list[Representation]
+):
+    """Add trait representations to instance.
+
+    Args:
+        instance (pyblish.api.Instance): Instance to add trait
+            representations to.
+        representations (list[Representation]): List of representation
+            trait based representations to add.
+
+    """
+    if not has_trait_representations(instance):
+        instance.data[TRAIT_INSTANCE_KEY] = []
+    instance.data[TRAIT_INSTANCE_KEY].extend(representations)
+
+
+def set_trait_representations(
+        instance: pyblish.api.Instance,
+        representations: list[Representation]
+):
+    """Set trait representations to instance.
+
+    Args:
+        instance (pyblish.api.Instance): Instance to set trait
+            representations to.
+        representations (list[Representation]): List of trait
+            based representations.
+
+    """
+    instance.data[TRAIT_INSTANCE_KEY] = representations
+
+
+def get_trait_representations(
+        instance: pyblish.api.Instance) -> list[Representation]:
+    """Get trait representations from instance.
+
+    Args:
+        instance (pyblish.api.Instance): Instance to get trait
+            representations from.
+
+    Returns:
+        list[Representation]: List of representation names.
+
+    """
+    return instance.data.get(TRAIT_INSTANCE_KEY, [])
diff --git a/client/ayon_core/plugins/publish/integrate_traits.py b/client/ayon_core/plugins/publish/integrate_traits.py
index 72a7ddd479..5587d64604 100644
--- a/client/ayon_core/plugins/publish/integrate_traits.py
+++ b/client/ayon_core/plugins/publish/integrate_traits.py
@@ -29,6 +29,9 @@ from ayon_core.lib.file_transaction import (
 from ayon_core.pipeline.publish import (
     PublishError,
     get_publish_template_name,
+    has_trait_representations,
+    get_trait_representations,
+    set_trait_representations,
 )
 from ayon_core.pipeline.traits import (
     UDIM,
@@ -236,17 +239,20 @@ class IntegrateTraits(pyblish.api.InstancePlugin):
             return
 
         # TODO (antirotor): Find better name for the key
-        if not instance.data.get("representations_with_traits"):
+        if not has_trait_representations(instance):
             self.log.debug(
                 "Instance has no representations with traits. Skipping")
             return
 
         # 2) filter representations based on LifeCycle traits
-        instance.data["representations_with_traits"] = self.filter_lifecycle(
-            instance.data["representations_with_traits"]
+        set_trait_representations(
+            instance,
+            self.filter_lifecycle(get_trait_representations(instance))
         )
 
-        representations: list[Representation] = instance.data["representations_with_traits"]  # noqa: E501
+        representations: list[Representation] = get_trait_representations(
+            instance
+        )
         if not representations:
             self.log.debug(
                 "Instance has no persistent representations. Skipping")