From 60f10feeeea54606dd278175237da2a27001aab1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20Samohel?= <ondrej.samohel@ynput.io>
Date: Fri, 29 Nov 2024 11:28:19 +0100
Subject: [PATCH] :art: add dict-like behavior to Representation

that and some tests
---
 .../pipeline/traits/representation.py         | 67 ++++++++++++++++++-
 .../ayon_core/pipeline/traits/test_traits.py  | 15 ++++-
 2 files changed, 80 insertions(+), 2 deletions(-)

diff --git a/client/ayon_core/pipeline/traits/representation.py b/client/ayon_core/pipeline/traits/representation.py
index 34aadd590a..756c4dfb67 100644
--- a/client/ayon_core/pipeline/traits/representation.py
+++ b/client/ayon_core/pipeline/traits/representation.py
@@ -1,6 +1,7 @@
 """Defines the base trait model and representation."""
 from __future__ import annotations
 
+import contextlib
 import inspect
 import re
 import sys
@@ -49,7 +50,7 @@ class Representation:
         representation_id (str): Representation ID.
 
     """
-    _data: dict
+    _data: dict[str, T]
     _module_blacklist: ClassVar[list[str]] = [
         "_", "builtins", "pydantic",
     ]
@@ -60,6 +61,70 @@ class Representation:
         """Return hash of the representation ID."""
         return hash(self.representation_id)
 
+    def __getitem__(self, key: str) -> T:
+        """Get the trait by ID.
+
+        Args:
+            key (str): Trait ID.
+
+        Returns:
+            TraitBase: Trait instance.
+
+        Raises:
+            MissingTraitError: If the trait is not found.
+
+        """
+        return self.get_trait_by_id(key)
+
+    def __setitem__(self, key: str, value: T) -> None:
+        """Set the trait by ID.
+
+        Args:
+            key (str): Trait ID.
+            value (TraitBase): Trait instance.
+
+        """
+        with contextlib.suppress(KeyError):
+            self._data.pop(key)
+
+        self.add_trait(value)
+
+    def __delitem__(self, key: str) -> None:
+        """Remove the trait by ID.
+
+        Args:
+            key (str): Trait ID.
+
+        Raises:
+            ValueError: If the trait is not found.
+
+        """
+        self.remove_trait_by_id(key)
+
+    def __contains__(self, key: str) -> bool:
+        """Check if the trait exists by ID.
+
+        Args:
+            key (str): Trait ID.
+
+        Returns:
+            bool: True if the trait exists, False otherwise.
+
+        """
+        return self.contains_trait_by_id(key)
+
+    def __iter__(self):
+        """Return the trait ID iterator."""
+        return iter(self._data)
+
+    def __str__(self):
+        """Return the representation name."""
+        return self.name
+
+    def items(self) -> dict[str, T]:
+        """Return the traits as items."""
+        return self._data.items()
+
     def add_trait(self, trait: TraitBase, *, exists_ok: bool=False) -> None:
         """Add a trait to the Representation.
 
diff --git a/tests/client/ayon_core/pipeline/traits/test_traits.py b/tests/client/ayon_core/pipeline/traits/test_traits.py
index 319401b91e..b8fa3962cc 100644
--- a/tests/client/ayon_core/pipeline/traits/test_traits.py
+++ b/tests/client/ayon_core/pipeline/traits/test_traits.py
@@ -25,15 +25,18 @@ REPRESENTATION_DATA = {
             "file_path": Path("/path/to/file"),
             "file_size": 1024,
             "file_hash": None,
+            "persitent": True,
         },
-        Image.id: {},
+        Image.id: {"persitent": True},
         PixelBased.id: {
             "display_window_width": 1920,
             "display_window_height": 1080,
             "pixel_aspect_ratio": 1.0,
+            "persitent": True,
         },
         Planar.id: {
             "planar_configuration": "RGB",
+            "persitent": True,
         },
     }
 
@@ -168,6 +171,16 @@ def test_trait_removing(representation: Representation) -> None:
             ValueError, match=f"Trait with ID {Image.id} not found."):
         representation.remove_trait(Image)
 
+def test_representation_dict_properties(representation: Representation) -> None:
+    """Test representation as dictionary."""
+    representation = Representation(name="test")
+    representation[Image.id] = Image()
+    assert Image.id in representation
+    image = representation[Image.id]
+    assert image == Image()
+    for trait_id, trait in representation.items():
+        assert trait_id == Image.id
+        assert trait == Image()
 
 
 def test_getting_traits_data(representation: Representation) -> None: