From 092325e64e09dd4cd60c97654ecf14224bca2e41 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ond=C5=99ej=20Samohel?= Date: Wed, 9 Oct 2024 17:40:22 +0200 Subject: [PATCH] :art: implement `TraitsData` --- client/ayon_core/pipeline/traits/__init__.py | 11 +- client/ayon_core/pipeline/traits/content.py | 89 ++++++++ .../pipeline/traits/three_dimensional.py | 8 +- client/ayon_core/pipeline/traits/trait.py | 210 ++++++++++++++++-- .../pipeline/traits/two_dimensional.py | 50 ++--- 5 files changed, 307 insertions(+), 61 deletions(-) create mode 100644 client/ayon_core/pipeline/traits/content.py diff --git a/client/ayon_core/pipeline/traits/__init__.py b/client/ayon_core/pipeline/traits/__init__.py index e85e88269b..1dbac8764d 100644 --- a/client/ayon_core/pipeline/traits/__init__.py +++ b/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", ] diff --git a/client/ayon_core/pipeline/traits/content.py b/client/ayon_core/pipeline/traits/content.py new file mode 100644 index 0000000000..1db36856d9 --- /dev/null +++ b/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") diff --git a/client/ayon_core/pipeline/traits/three_dimensional.py b/client/ayon_core/pipeline/traits/three_dimensional.py index c27b7ee56b..0638700ab7 100644 --- a/client/ayon_core/pipeline/traits/three_dimensional.py +++ b/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") diff --git a/client/ayon_core/pipeline/traits/trait.py b/client/ayon_core/pipeline/traits/trait.py index 0ac2d63cdd..c6b7258535 100644 --- a/client/ayon_core/pipeline/traits/trait.py +++ b/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) diff --git a/client/ayon_core/pipeline/traits/two_dimensional.py b/client/ayon_core/pipeline/traits/two_dimensional.py index cd180ffc6b..0b6fea315b 100644 --- a/client/ayon_core/pipeline/traits/two_dimensional.py +++ b/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")