♻️ wrap trait based instance into functions

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",
 )

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, [])

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