🎨 implement TraitsData

client/ayon_core/pipeline/traits/__init__.py

@@ -1,7 +1,8 @@
 """Trait classes for the pipeline."""
-from .trait import TraitBase
+from .content import Compressed, FileLocation, RootlessLocation
+from .three_dimensional import Spatial
+from .trait import TraitBase, TraitsData
 from .two_dimensional import (
-    Compressed,
     Deep,
     Image,
     Overscan,
@@ -12,6 +13,10 @@ from .two_dimensional import (
 __all__ = [
     # base
     "TraitBase",
+    "TraitsData",
+    # content
+    "FileLocation",
+    "RootlessLocation",
     # two-dimensional
     "Image",
     "PixelBased",
@@ -19,4 +24,6 @@ __all__ = [
     "Deep",
     "Compressed",
     "Overscan",
+    # three-dimensional
+    "Spatial",
 ]

client/ayon_core/pipeline/traits/content.py

@@ -0,0 +1,89 @@
+"""Content traits for the pipeline."""
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, ClassVar, Optional
+
+from pydantic import Field
+
+from .trait import TraitBase
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+
+class MimeType(TraitBase):
+    """MimeType trait model.
+
+    This model represents a mime type trait.
+
+    Attributes:
+        name (str): Trait name.
+        description (str): Trait description.
+        id (str): id should be namespaced trait name with version
+        mime_type (str): Mime type.
+
+    """
+
+    name: ClassVar[str] = "MimeType"
+    description: ClassVar[str] = "MimeType Trait Model"
+    id: ClassVar[str] = "ayon.content.MimeType.v1"
+    mime_type: str = Field(..., title="Mime Type")
+
+class FileLocation(TraitBase):
+    """FileLocation trait model.
+
+    This model represents a file location trait.
+
+    Attributes:
+        name (str): Trait name.
+        description (str): Trait description.
+        id (str): id should be namespaced trait name with version
+        file_path (str): File path.
+        file_size (int): File size in bytes.
+        file_hash (str): File hash.
+
+    """
+
+    name: ClassVar[str] = "FileLocation"
+    description: ClassVar[str] = "FileLocation Trait Model"
+    id: ClassVar[str] = "ayon.content.FileLocation.v1"
+    file_path: Path = Field(..., title="File Path")
+    file_size: int = Field(..., title="File Size")
+    file_hash: Optional[str] = Field(..., title="File Hash")
+
+class RootlessLocation(TraitBase):
+    """RootlessLocation trait model.
+
+    This model represents a rootless location trait.
+
+    Attributes:
+        name (str): Trait name.
+        description (str): Trait description.
+        id (str): id should be namespaced trait name with version
+        rootless_path (str): Rootless path.
+
+    """
+
+    name: ClassVar[str] = "RootlessLocation"
+    description: ClassVar[str] = "RootlessLocation Trait Model"
+    id: ClassVar[str] = "ayon.content.RootlessLocation.v1"
+    rootless_path: str = Field(..., title="File Path")
+
+
+class Compressed(TraitBase):
+    """Compressed trait model.
+
+    This model represents a compressed trait.
+
+    Attributes:
+        name (str): Trait name.
+        description (str): Trait description.
+        id (str): id should be namespaced trait name with version
+        compression_type (str): Compression type.
+
+    """
+
+    name: ClassVar[str] = "Compressed"
+    description: ClassVar[str] = "Compressed Trait"
+    id: ClassVar[str] = "ayon.content.Compressed.v1"
+    compression_type: str = Field(..., title="Compression Type")

client/ayon_core/pipeline/traits/three_dimensional.py

@@ -1,4 +1,6 @@
 """Two-dimensional image traits."""
+from typing import ClassVar
+
 from pydantic import Field
 
 from .trait import TraitBase
@@ -13,9 +15,9 @@ class Spatial(TraitBase):
         meters_per_unit (float): Meters per unit.
 
     """
-    id: str = "ayon.content.Spatial.v1"
-    name: str = "Spatial"
-    description = "Spatial trait model."
+    id: ClassVar[str] = "ayon.3d.Spatial.v1"
+    name: ClassVar[str] = "Spatial"
+    description: ClassVar[str] = "Spatial trait model."
     up_axis: str = Field(..., title="Up axis")
     handedness: str = Field(..., title="Handedness")
     meters_per_unit: float = Field(..., title="Meters per unit")

client/ayon_core/pipeline/traits/trait.py

@@ -1,35 +1,197 @@
 """Defines the base trait model."""
-from pydantic import BaseModel, Field
+from __future__ import annotations
+
+import inspect
+import sys
+from abc import ABC, abstractmethod
+from functools import lru_cache
+from typing import ClassVar, Optional, Type, Union
+
+import pydantic.alias_generators
+from pydantic import AliasGenerator, BaseModel, ConfigDict
 
 
-def camelize(src: str) -> str:
-    """Convert snake_case to camelCase."""
-    components = src.split("_")
-    return components[0] + "".join(x.title() for x in components[1:])
-
-
-class TraitBase(BaseModel):
+class TraitBase(ABC, BaseModel):
     """Base trait model.
 
     This model must be used as a base for all trait models.
 
-    Attributes:
-        name (str): Trait name.
-        description (str): Trait description.
-        id (str): id should be namespaced trait name with version
-            as ``ayon.content.LocatableBundle.v1``
-
     """
 
-    class Config:
-        """API model config."""
+    model_config = ConfigDict(
+        alias_generator=AliasGenerator(
+            serialization_alias=pydantic.alias_generators.to_camel,
+        )
+    )
 
-        orm_mode = True
-        allow_population_by_field_name = True
-        alias_generator = camelize
+    @property
+    @abstractmethod
+    def id(self) -> str:
+        """Abstract attribute for ID."""
+        ...
 
-    name: str = Field(..., title="Trait name")
-    description: str = Field(..., title="Trait description")
-    # id should be: ayon.content.LocatableBundle.v1
-    id: str = Field(..., title="Trait ID",
-                    description="Unique identifier for the trait.")
+    @property
+    @abstractmethod
+    def name(self) -> str:
+        """Abstract attribute for name."""
+        ...
+
+    @property
+    @abstractmethod
+    def description(self) -> str:
+        """Abstract attribute for description."""
+        ...
+
+
+
+class TraitsData:
+    """Traits data container.
+
+    This model represents the data of a trait.
+
+    """
+    _data: dict
+    _module_blacklist: ClassVar[list[str]] = [
+        "_", "builtins", "pydantic"]
+
+    @lru_cache(maxsize=64)  # noqa: B019
+    def _get_trait_class(self, trait_id: str) -> Union[Type[TraitBase], None]:
+        """Get the trait class with corresponding to given ID.
+
+        This method will search for the trait class in all the modules except
+        the blacklisted modules. There is some issue in Pydantic where
+        ``issubclass`` is not working properly so we are excluding explicitly
+        modules with offending classes. This list can be updated as needed to
+        speed up the search.
+
+        Args:
+            trait_id (str): Trait ID.
+
+        Returns:
+            Type[TraitBase]: Trait class.
+
+        """
+        modules = sys.modules.copy()
+        filtered_modules = modules.copy()
+        for module_name in modules:
+            for bl_module in self._module_blacklist:
+                if module_name.startswith(bl_module):
+                    filtered_modules.pop(module_name)
+
+        for module in filtered_modules.values():
+            if not module:
+                continue
+            for _, klass in inspect.getmembers(module, inspect.isclass):
+                if inspect.isclass(klass) and \
+                        issubclass(klass, TraitBase) and \
+                        klass.id == trait_id:
+                    return klass
+        return None
+
+
+    def add(self, trait: TraitBase, *, exists_ok: bool=False) -> None:
+        """Add a trait to the data.
+
+        Args:
+            trait (TraitBase): Trait to add.
+            exists_ok (bool, optional): If True, do not raise an error if the
+                trait already exists. Defaults to False.
+
+        Raises:
+            ValueError: If the trait ID is not provided or the trait already
+                exists.
+
+        """
+        if not trait.id:
+            error_msg = f"Invalid trait {trait} - ID is required."
+            raise ValueError(error_msg)
+        if trait.id in self._data and not exists_ok:
+            error_msg = f"Trait with ID {trait.id} already exists."
+            raise ValueError(error_msg)
+        self._data[trait.id] = trait
+
+    def remove(self,
+               trait_id: Optional[str],
+               trait: Optional[Type[TraitBase]]) -> None:
+        """Remove a trait from the data.
+
+        Args:
+            trait_id (str, optional): Trait ID.
+            trait (TraitBase, optional): Trait class.
+
+        """
+        if trait_id:
+            self._data.pop(trait_id)
+        elif trait:
+            self._data.pop(trait.id)
+
+    def has_trait(self,
+                  trait_id: Optional[str]=None,
+                  trait: Optional[Type[TraitBase]]=None) -> bool:
+        """Check if the trait exists.
+
+        Args:
+            trait_id (str, optional): Trait ID.
+            trait (TraitBase, optional): Trait class.
+
+        Returns:
+            bool: True if the trait exists, False otherwise.
+
+        """
+        if not trait_id:
+            trait_id = trait.id
+        return hasattr(self, trait_id)
+
+    def get(self,
+            trait_id: Optional[str]=None,
+            trait: Optional[Type[TraitBase]]=None) -> Union[TraitBase, None]:
+        """Get a trait from the data.
+
+        Args:
+            trait_id (str, optional): Trait ID.
+            trait (TraitBase, optional): Trait class.
+
+        Returns:
+            TraitBase: Trait instance.
+
+        """
+        trait_class = None
+        if trait_id:
+            trait_class = self._get_trait_class(trait_id)
+            if not trait_class:
+                error_msg = f"Trait model with ID {trait_id} not found."
+                raise ValueError(error_msg)
+
+        if trait:
+            trait_class = trait
+            trait_id = trait.id
+
+        if not trait_class and not trait_id:
+            error_msg = "Trait ID or Trait class is required"
+            raise ValueError(error_msg)
+
+        return self._data[trait_id] if self._data.get(trait_id) else None
+
+    def as_dict(self) -> dict:
+        """Return the data as a dictionary.
+
+        Returns:
+            dict: Data dictionary.
+
+        """
+        result = {
+            trait_id: dict(sorted(trait.dict()))
+            for trait_id, trait in self._data.items()
+        }
+        return dict(sorted(result))
+
+    def __len__(self):
+        """Return the length of the data."""
+        return len(self._data)
+
+    def __init__(self, traits: Optional[list[TraitBase]]):
+        """Initialize the data."""
+        self._data = {}
+        if traits:
+            for trait in traits:
+                self.add(trait)

client/ayon_core/pipeline/traits/two_dimensional.py

@@ -1,4 +1,6 @@
 """Two-dimensional image traits."""
+from typing import ClassVar
+
 from pydantic import Field
 
 from .trait import TraitBase
@@ -16,9 +18,9 @@ class Image(TraitBase):
 
     """
 
-    name: str = "Image"
-    description = "Image Trait"
-    id: str = "ayon.content.Image.v1"
+    name: ClassVar[str] = "Image"
+    description: ClassVar[str] = "Image Trait"
+    id: ClassVar[str] = "ayon.2d.Image.v1"
 
 
 class PixelBased(TraitBase):
@@ -36,9 +38,9 @@ class PixelBased(TraitBase):
 
     """
 
-    name: str = "PixelBased"
-    description = "PixelBased Trait Model"
-    id: str = "ayon.content.PixelBased.v1"
+    name: ClassVar[str] = "PixelBased"
+    description: ClassVar[str] = "PixelBased Trait Model"
+    id: ClassVar[str] = "ayon.2d.PixelBased.v1"
     display_window_width: int = Field(..., title="Display Window Width")
     display_window_height: int = Field(..., title="Display Window Height")
     pixel_aspect_ratio: float = Field(..., title="Pixel Aspect Ratio")
@@ -49,7 +51,7 @@ class Planar(TraitBase):
 
     This model represents an Image with planar configuration.
 
-    Todo (antirotor): Is this really a planar configuration? As with
+    TODO (antirotor): Is this really a planar configuration? As with
         bitplanes and everything? If it serves as differentiator for
         Deep images, should it be named differently? Like Raster?
 
@@ -61,9 +63,9 @@ class Planar(TraitBase):
 
     """
 
-    name: str = "Planar"
-    description = "Planar Trait Model"
-    id: str = "ayon.content.Planar.v1"
+    name: ClassVar[str] = "Planar"
+    description: ClassVar[str] = "Planar Trait Model"
+    id: ClassVar[str] = "ayon.2d.Planar.v1"
     planar_configuration: str = Field(..., title="Planar-based Image")
 
 
@@ -80,29 +82,13 @@ class Deep(TraitBase):
 
     """
 
-    name: str = "Deep"
-    description = "Deep Trait Model"
-    id: str = "ayon.content.Deep.v1"
+    name: ClassVar[str] = "Deep"
+    description: ClassVar[str] = "Deep Trait Model"
+    id: ClassVar[str] = "ayon.2d.Deep.v1"
     deep_data_type: str = Field(..., title="Deep Data Type")
 
 
-class Compressed(TraitBase):
-    """Compressed trait model.
 
-    This model represents a compressed image trait.
-
-    Attributes:
-        name (str): Trait name.
-        description (str): Trait description.
-        id (str): id should be namespaced trait name with version
-        compression_type (str): Compression type.
-
-    """
-
-    name: str = "Compressed"
-    description = "Compressed Trait"
-    id: str = "ayon.content.Compressed.v1"
-    compression_type: str = Field(..., title="Compression Type")
 
 
 class Overscan(TraitBase):
@@ -121,9 +107,9 @@ class Overscan(TraitBase):
 
     """
 
-    name: str = "Overscan"
-    description = "Overscan Trait"
-    id: str = "ayon.content.Overscan.v1"
+    name: ClassVar[str] = "Overscan"
+    description: ClassVar[str] = "Overscan Trait"
+    id: ClassVar[str] = "ayon.2d.Overscan.v1"
     left: int = Field(..., title="Left Overscan")
     right: int = Field(..., title="Right Overscan")
     top: int = Field(..., title="Top Overscan")