From e0a6e1767c6cd827d425a97ad612a7ae7b968af6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ond=C5=99ej=20Samohel?= Date: Mon, 14 Oct 2024 22:36:16 +0200 Subject: [PATCH] :recycle: add traits and refactor api --- client/ayon_core/pipeline/traits/__init__.py | 29 ++- client/ayon_core/pipeline/traits/content.py | 29 ++- client/ayon_core/pipeline/traits/time.py | 81 +++++++ client/ayon_core/pipeline/traits/trait.py | 215 ++++++++++++------ .../ayon_core/pipeline/traits/test_traits.py | 164 ++++++++++--- 5 files changed, 410 insertions(+), 108 deletions(-) create mode 100644 client/ayon_core/pipeline/traits/time.py diff --git a/client/ayon_core/pipeline/traits/__init__.py b/client/ayon_core/pipeline/traits/__init__.py index 429b307e3d..2d17cccf7e 100644 --- a/client/ayon_core/pipeline/traits/__init__.py +++ b/client/ayon_core/pipeline/traits/__init__.py @@ -1,6 +1,13 @@ """Trait classes for the pipeline.""" -from .content import Compressed, FileLocation, RootlessLocation +from .content import ( + Bundle, + Compressed, + FileLocation, + MimeType, + RootlessLocation, +) from .three_dimensional import Spatial +from .time import Clip, GapPolicy, Sequence, SMPTETimecode from .trait import Representation, TraitBase from .two_dimensional import ( Deep, @@ -12,18 +19,30 @@ from .two_dimensional import ( __all__ = [ # base - "TraitBase", "Representation", + "TraitBase", + # content + "Bundle", + "Compressed", "FileLocation", + "MimeType", "RootlessLocation", + # two-dimensional + "Compressed", + "Deep", "Image", + "Overscan", "PixelBased", "Planar", - "Deep", - "Compressed", - "Overscan", + # three-dimensional "Spatial", + + # time + "Clip", + "GapPolicy", + "Sequence", + "SMPTETimecode", ] diff --git a/client/ayon_core/pipeline/traits/content.py b/client/ayon_core/pipeline/traits/content.py index 3eec848f69..1cb7366928 100644 --- a/client/ayon_core/pipeline/traits/content.py +++ b/client/ayon_core/pipeline/traits/content.py @@ -7,7 +7,7 @@ from typing import ClassVar, Optional from pydantic import Field -from .trait import TraitBase +from .trait import Representation, TraitBase class MimeType(TraitBase): @@ -86,3 +86,30 @@ class Compressed(TraitBase): description: ClassVar[str] = "Compressed Trait" id: ClassVar[str] = "ayon.content.Compressed.v1" compression_type: str = Field(..., title="Compression Type") + + +class Bundle(TraitBase): + """Bundle trait model. + + This model list of independent Representation traits + that are bundled together. This is useful for representing + a collection of representations that are part of a single + entity. + + Attributes: + name (str): Trait name. + description (str): Trait description. + id (str): id should be namespaced trait name with version + items (list[list[TraitBase]]): List of representations. + + """ + + name: ClassVar[str] = "Bundle" + description: ClassVar[str] = "Bundle Trait" + id: ClassVar[str] = "ayon.content.Bundle.v1" + items: list[list[TraitBase]] = Field( + ..., title="Bundles of traits") + + def to_representation(self) -> Representation: + """Convert to a representation.""" + return Representation(traits=self.items) diff --git a/client/ayon_core/pipeline/traits/time.py b/client/ayon_core/pipeline/traits/time.py new file mode 100644 index 0000000000..de4ff53790 --- /dev/null +++ b/client/ayon_core/pipeline/traits/time.py @@ -0,0 +1,81 @@ +"""Temporal (time related) traits.""" +from enum import Enum, auto +from typing import ClassVar + +from pydantic import Field + +from .trait import TraitBase + + +class GapPolicy(Enum): + """Gap policy enumeration. + + Attributes: + forbidden (int): Gaps are forbidden. + missing (int): Gaps are interpreted as missing frames. + hold (int): Gaps are interpreted as hold frames (last existing frames). + black (int): Gaps are interpreted as black frames. + """ + forbidden = auto() + missing = auto() + hold = auto() + black = auto() + +class Clip(TraitBase): + """Clip trait model. + + Model representing a clip trait. + + Attributes: + name (str): Trait name. + description (str): Trait description. + id (str): id should be namespaced trait name with version + frame_start (int): Frame start. + frame_end (int): Frame end. + frame_start_handle (int): Frame start handle. + frame_end_handle (int): Frame end handle. + + """ + name: ClassVar[str] = "Clip" + description: ClassVar[str] = "Clip Trait" + id: ClassVar[str] = "ayon.time.Clip.v1" + frame_start: int = Field(..., title="Frame Start") + frame_end: int = Field(..., title="Frame End") + frame_start_handle: int = Field(..., title="Frame Start Handle") + frame_end_handle: int = Field(..., title="Frame End Handle") + +class Sequence(Clip): + """Sequence trait model. + + This model represents a sequence trait. Based on the Clip trait, + adding handling for steps, gaps policy and frame padding. + + Attributes: + name (str): Trait name. + description (str): Trait description. + id (str): id should be namespaced trait name with version + step (int): Frame step. + gaps_policy (GapPolicy): Gaps policy - how to handle gaps in + sequence. + frame_padding (int): Frame padding. + frame_regex (str): Frame regex - regular expression to match + frame numbers. + + """ + name: ClassVar[str] = "Sequence" + description: ClassVar[str] = "Sequence Trait Model" + id: ClassVar[str] = "ayon.time.Sequence.v1" + step: int = Field(..., title="Step") + gaps_policy: GapPolicy = Field( + GapPolicy.forbidden, title="Gaps Policy") + frame_padding: int = Field(..., title="Frame Padding") + frame_regex: str = Field(..., title="Frame Regex") + + +# Do we need one for drop and non-drop frame? +class SMPTETimecode(TraitBase): + """Timecode trait model.""" + name: ClassVar[str] = "Timecode" + description: ClassVar[str] = "SMPTE Timecode Trait" + id: ClassVar[str] = "ayon.time.SMPTETimecode.v1" + timecode: str = Field(..., title="SMPTE Timecode HH:MM:SS:FF") diff --git a/client/ayon_core/pipeline/traits/trait.py b/client/ayon_core/pipeline/traits/trait.py index fc71b3d5d0..6fa36aeb51 100644 --- a/client/ayon_core/pipeline/traits/trait.py +++ b/client/ayon_core/pipeline/traits/trait.py @@ -132,130 +132,197 @@ class Representation: for trait in traits: self.add_trait(trait, exists_ok=exists_ok) - def remove_trait(self, - trait_id: Optional[str]=None, - trait: Optional[Type[TraitBase]]=None) -> None: + def remove_trait(self, trait: 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) + Raises: + ValueError: If the trait is not found. - def remove_traits(self, - trait_ids: Optional[list[str]]=None, - traits: Optional[list[Type[TraitBase]]]=None) -> None: + """ + try: + self._data.pop(trait.id) + except KeyError as e: + error_msg = f"Trait with ID {trait.id} not found." + raise ValueError(error_msg) from e + + def remove_trait_by_id(self, trait_id: str) -> None: + """Remove a trait from the data by its ID. + + Args: + trait_id (str): Trait ID. + + Raises: + ValueError: If the trait is not found. + + """ + try: + self._data.pop(trait_id) + except KeyError as e: + error_msg = f"Trait with ID {trait_id} not found." + raise ValueError(error_msg) from e + + def remove_traits(self, traits: list[Type[TraitBase]]) -> None: """Remove a list of traits from the Representation. + If no trait IDs or traits are provided, all traits will be removed. + + Args: + traits (list[TraitBase]): List of trait classes. + + """ + if not traits: + self._data = {} + return + + for trait in traits: + self.remove_trait(trait) + + def remove_traits_by_id(self, trait_ids: list[str]) -> None: + """Remove a list of traits from the Representation by their ID. + + If no trait IDs or traits are provided, all traits will be removed. + Args: trait_ids (list[str], optional): List of trait IDs. - traits (list[TraitBase], optional): List of trait classes. """ - if trait_ids: - for trait_id in trait_ids: - self.remove_trait(trait_id=trait_id) - elif traits: - for trait in traits: - self.remove_trait(trait=trait) + for trait_id in trait_ids: + self.remove_trait_by_id(trait_id) - def has_trait(self, - trait_id: Optional[str]=None, - trait: Optional[Type[TraitBase]]=None) -> bool: - """Check if the trait exists. + + def has_traits(self) -> bool: + """Check if the Representation has any traits. + + Returns: + bool: True if the Representation has any traits, False otherwise. + + """ + return bool(self._data) + + def contains_trait(self, trait: Type[TraitBase]) -> bool: + """Check if the trait exists in the Representation. Args: - trait_id (str, optional): Trait ID. - trait (TraitBase, optional): Trait class. + trait (TraitBase): Trait class. Returns: bool: True if the trait exists, False otherwise. """ - if not trait_id: - trait_id = trait.id - return hasattr(self, trait_id) + return bool(self._data.get(trait.id)) - def has_traits(self, - trait_ids: Optional[list[str]]=None, - traits: Optional[list[Type[TraitBase]]]=None) -> bool: + def contains_trait_by_id(self, trait_id: str) -> bool: + """Check if the trait exists using trait id. + + Args: + trait_id (str): Trait ID. + + Returns: + bool: True if the trait exists, False otherwise. + + """ + return bool(self._data.get(trait_id)) + + def contains_traits(self, traits: list[Type[TraitBase]]) -> bool: """Check if the traits exist. Args: - trait_ids (list[str], optional): List of trait IDs. traits (list[TraitBase], optional): List of trait classes. Returns: bool: True if all traits exist, False otherwise. """ - if trait_ids: - for trait_id in trait_ids: - if not self.has_trait(trait_id=trait_id): - return False - elif traits: - for trait in traits: - if not self.has_trait(trait=trait): - return False - return True + return all(self.contains_trait(trait=trait) for trait in traits) - def get_trait(self, - trait_id: Optional[str]=None, - trait: Optional[Type[TraitBase]]=None - ) -> Union[TraitBase, None]: + def contains_traits_by_id(self, trait_ids: list[str]) -> bool: + """Check if the traits exist by id. + + If no trait IDs or traits are provided, it will check if the + representation has any traits. + + Args: + trait_ids (list[str]): List of trait IDs. + + Returns: + bool: True if all traits exist, False otherwise. + + """ + return all( + self.contains_trait_by_id(trait_id) for trait_id in trait_ids + ) + + def get_trait(self, trait: Type[TraitBase]) -> Union[TraitBase, None]: """Get a trait from the representation. 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) + return self._data[trait.id] if self._data.get(trait.id) else None - if trait: - trait_class = trait - trait_id = trait.id + def get_trait_by_id(self, trait_id: str) -> Union[TraitBase, None]: + """Get a trait from the representation by id. - if not trait_class and not trait_id: - error_msg = "Trait ID or Trait class is required" + Args: + trait_id (str): Trait ID. + + Returns: + TraitBase: Trait instance. + + """ + 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) return self._data[trait_id] if self._data.get(trait_id) else None def get_traits(self, - trait_ids: Optional[list[str]]=None, traits: Optional[list[Type[TraitBase]]]=None) -> dict: - """Get a list of traits from the representation. + """Get a list of traits from the representation. - Args: - trait_ids (list[str], optional): List of trait IDs. - traits (list[TraitBase], optional): List of trait classes. + If no trait IDs or traits are provided, all traits will be returned. - Returns: - dict: Dictionary of traits. + Args: + traits (list[TraitBase], optional): List of trait classes. - """ - result = {} - if trait_ids: - for trait_id in trait_ids: - result[trait_id] = self.get_trait(trait_id=trait_id) - elif traits: - for trait in traits: - result[trait.id] = self.get_trait(trait=trait) - return result + Returns: + dict: Dictionary of traits. + + """ + result = {} + if not traits: + for trait_id in self._data: + result[trait_id] = self.get_trait_by_id(trait_id=trait_id) + return result + + for trait in traits: + result[trait.id] = self.get_trait(trait=trait) + return result + + def get_traits_by_ids(self, trait_ids: list[str]) -> dict: + """Get a list of traits from the representation by their id. + + If no trait IDs or traits are provided, all traits will be returned. + + Args: + trait_ids (list[str]): List of trait IDs. + + Returns: + dict: Dictionary of traits. + + """ + return { + trait_id: self.get_trait_by_id(trait_id) + for trait_id in trait_ids + } def traits_as_dict(self) -> dict: """Return the traits from Representation data as a dictionary. @@ -276,7 +343,7 @@ class Representation: """Return the length of the data.""" return len(self._data) - def __init__(self, traits: Optional[list[TraitBase]]): + def __init__(self, traits: Optional[list[TraitBase]]=None): """Initialize the data.""" self._data = {} if traits: diff --git a/tests/client/ayon_core/pipeline/traits/test_traits.py b/tests/client/ayon_core/pipeline/traits/test_traits.py index 0bf08219ed..e6ca0a1067 100644 --- a/tests/client/ayon_core/pipeline/traits/test_traits.py +++ b/tests/client/ayon_core/pipeline/traits/test_traits.py @@ -5,8 +5,10 @@ from pathlib import Path import pytest from ayon_core.pipeline.traits import ( + Bundle, FileLocation, Image, + MimeType, PixelBased, Planar, Representation, @@ -41,49 +43,111 @@ def representation() -> Representation: Planar(**REPRESENTATION_DATA[Planar.id]), ]) +def test_representation_errors(representation: Representation) -> None: + """Test errors in representation.""" + with pytest.raises(ValueError, + match="Trait ID or Trait class is required"): + representation.get_trait() + + with pytest.raises(ValueError, + match="Trait ID or Trait class is required"): + representation.contains_trait() + def test_representation_traits(representation: Representation) -> None: """Test setting and getting traits.""" assert len(representation) == len(REPRESENTATION_DATA) - assert representation.get_trait(trait_id=FileLocation.id) - assert representation.get_trait(trait_id=Image.id) - assert representation.get_trait(trait_id=PixelBased.id) - assert representation.get_trait(trait_id=Planar.id) + assert representation.get_trait_by_id(FileLocation.id) + assert representation.get_trait_by_id(Image.id) + assert representation.get_trait_by_id(trait_id="ayon.2d.Image.v1") + assert representation.get_trait_by_id(PixelBased.id) + assert representation.get_trait_by_id(trait_id="ayon.2d.PixelBased.v1") + assert representation.get_trait_by_id(Planar.id) + assert representation.get_trait_by_id(trait_id="ayon.2d.Planar.v1") - assert representation.get_trait(trait=FileLocation) - assert representation.get_trait(trait=Image) - assert representation.get_trait(trait=PixelBased) - assert representation.get_trait(trait=Planar) + assert representation.get_trait(FileLocation) + assert representation.get_trait(Image) + assert representation.get_trait(PixelBased) + assert representation.get_trait(Planar) assert issubclass( - type(representation.get_trait(trait=FileLocation)), TraitBase) + type(representation.get_trait(FileLocation)), TraitBase) - assert representation.get_trait( - trait=FileLocation) == representation.get_trait( - trait_id=FileLocation.id) - assert representation.get_trait( - trait=Image) == representation.get_trait( - trait_id=Image.id) - assert representation.get_trait( - trait=PixelBased) == representation.get_trait( - trait_id=PixelBased.id) - assert representation.get_trait( - trait=Planar) == representation.get_trait( - trait_id=Planar.id) + assert representation.get_trait(FileLocation) == \ + representation.get_trait_by_id(FileLocation.id) + assert representation.get_trait(Image) == \ + representation.get_trait_by_id(Image.id) + assert representation.get_trait(PixelBased) == \ + representation.get_trait_by_id(PixelBased.id) + assert representation.get_trait(Planar) == \ + representation.get_trait_by_id(Planar.id) - assert representation.get_trait(trait_id="ayon.2d.Image.v1") - assert representation.get_trait(trait_id="ayon.2d.PixelBased.v1") - assert representation.get_trait(trait_id="ayon.2d.Planar.v1") - - assert representation.get_trait( - trait_id="ayon.2d.PixelBased.v1").display_window_width == \ + assert representation.get_trait_by_id( + "ayon.2d.PixelBased.v1").display_window_width == \ REPRESENTATION_DATA[PixelBased.id]["display_window_width"] assert representation.get_trait( trait=PixelBased).display_window_height == \ REPRESENTATION_DATA[PixelBased.id]["display_window_height"] + repre_dict = { + FileLocation.id: FileLocation(**REPRESENTATION_DATA[FileLocation.id]), + Image.id: Image(), + PixelBased.id: PixelBased(**REPRESENTATION_DATA[PixelBased.id]), + Planar.id: Planar(**REPRESENTATION_DATA[Planar.id]), + } + assert representation.get_traits() == repre_dict + + assert representation.get_traits_by_ids( + trait_ids=[FileLocation.id, Image.id, PixelBased.id, Planar.id]) == \ + repre_dict + assert representation.get_traits( + [FileLocation, Image, PixelBased, Planar]) == \ + repre_dict + + assert representation.has_traits() is True + empty_representation = Representation(traits=[]) + assert empty_representation.has_traits() is False + + assert representation.contains_trait(trait=FileLocation) is True + assert representation.contains_traits([Image, FileLocation]) is True + assert representation.contains_trait_by_id(FileLocation.id) is True + assert representation.contains_traits_by_id( + trait_ids=[FileLocation.id, Image.id]) is True + + assert representation.contains_trait(trait=Bundle) is False + assert representation.contains_traits([Image, Bundle]) is False + assert representation.contains_trait_by_id(Bundle.id) is False + assert representation.contains_traits_by_id( + trait_ids=[FileLocation.id, Bundle.id]) is False + +def test_trait_removing(representation: Representation) -> None: + """Test removing traits.""" + assert representation.contains_trait_by_id("nonexistent") is False + with pytest.raises( + ValueError, match="Trait with ID nonexistent not found."): + representation.remove_trait_by_id("nonexistent") + + assert representation.contains_trait(trait=FileLocation) is True + representation.remove_trait(trait=FileLocation) + assert representation.contains_trait(trait=FileLocation) is False + + assert representation.contains_trait_by_id(Image.id) is True + representation.remove_trait_by_id(Image.id) + assert representation.contains_trait_by_id(Image.id) is False + + assert representation.contains_traits([PixelBased, Planar]) is True + representation.remove_traits([Planar, PixelBased]) + assert representation.contains_traits([PixelBased, Planar]) is False + + assert representation.has_traits() is False + + with pytest.raises( + ValueError, match=f"Trait with ID {Image.id} not found."): + representation.remove_trait(Image) + + def test_getting_traits_data(representation: Representation) -> None: """Test getting a batch of traits.""" - result = representation.get_traits( + result = representation.get_traits_by_ids( trait_ids=[FileLocation.id, Image.id, PixelBased.id, Planar.id]) assert result == { "ayon.2d.Image.v1": Image(), @@ -103,3 +167,47 @@ def test_traits_data_to_dict(representation: Representation) -> None: """Test converting traits data to dictionary.""" result = representation.traits_as_dict() assert result == REPRESENTATION_DATA + + +def test_bundles() -> None: + """Test bundle trait.""" + diffuse_texture = [ + Image(), + PixelBased( + display_window_width=1920, + display_window_height=1080, + pixel_aspect_ratio=1.0), + Planar(planar_configuration="RGB"), + FileLocation( + file_path=Path("/path/to/diffuse.jpg"), + file_size=1024, + file_hash=None), + MimeType(mime_type="image/jpeg"), + ] + bump_texture = [ + Image(), + PixelBased( + display_window_width=1920, + display_window_height=1080, + pixel_aspect_ratio=1.0), + Planar(planar_configuration="RGB"), + FileLocation( + file_path=Path("/path/to/bump.tif"), + file_size=1024, + file_hash=None), + MimeType(mime_type="image/tiff"), + ] + bundle = Bundle(items=[diffuse_texture, bump_texture]) + representation = Representation(traits=[bundle]) + + if representation.contains_trait(trait=Bundle): + assert representation.get_trait(trait=Bundle).items == [ + diffuse_texture, bump_texture + ] + + for item in representation.get_trait(trait=Bundle).items: + sub_representation = Representation(traits=item) + assert sub_representation.contains_trait(trait=Image) + assert sub_representation.get_trait(trait=MimeType).mime_type in [ + "image/jpeg", "image/tiff" + ]