diff --git a/client/ayon_core/pipeline/traits/two_dimensional.py b/client/ayon_core/pipeline/traits/two_dimensional.py index 93d21a9bc3..2cd27340ad 100644 --- a/client/ayon_core/pipeline/traits/two_dimensional.py +++ b/client/ayon_core/pipeline/traits/two_dimensional.py @@ -1,10 +1,15 @@ """Two-dimensional image traits.""" -from typing import ClassVar +from __future__ import annotations -from pydantic import Field +import re +from typing import TYPE_CHECKING, ClassVar, Optional + +from pydantic import Field, field_validator from .trait import TraitBase +if TYPE_CHECKING: + from .content import FileLocation, FileLocations class Image(TraitBase): """Image trait model. @@ -129,4 +134,39 @@ class UDIM(TraitBase): name: ClassVar[str] = "UDIM" description: ClassVar[str] = "UDIM Trait" id: ClassVar[str] = "ayon.2d.UDIM.v1" - udim: int = Field(..., title="UDIM") + udim: list[int] = Field(..., title="UDIM") + udim_regex: Optional[str] = Field( + r"(?:\.|_)(?P\d+)\.\D+\d?$", title="UDIM Regex") + + @field_validator("udim_regex") + @classmethod + def validate_frame_regex(cls, v: Optional[str]) -> str: + """Validate udim regex.""" + if v is not None and "?P" not in v: + msg = "UDIM regex must include 'udim' named group" + raise ValueError(msg) + return v + + def get_file_location_for_udim( + self, + file_locations: FileLocations, + udim: int, + ) -> Optional[FileLocation]: + """Get file location for UDIM. + + Args: + file_locations (FileLocations): File locations. + udim (int): UDIM value. + + Returns: + Optional[FileLocation]: File location. + + """ + pattern = re.compile(self.udim_regex) + for location in file_locations.file_paths: + result = re.search(pattern, location.file_path.name) + if result: + udim_index = int(result.group("udim")) + if udim_index == udim: + return location + return None diff --git a/tests/client/ayon_core/pipeline/traits/test_two_dimesional_traits.py b/tests/client/ayon_core/pipeline/traits/test_two_dimesional_traits.py new file mode 100644 index 0000000000..9428d86a39 --- /dev/null +++ b/tests/client/ayon_core/pipeline/traits/test_two_dimesional_traits.py @@ -0,0 +1,43 @@ +"""Tests for the 2d related traits.""" +from __future__ import annotations + +from pathlib import Path + +from ayon_core.pipeline.traits import ( + UDIM, + FileLocation, + FileLocations, + Representation, +) + + +def test_get_file_location_for_udim() -> None: + """Test get_file_location_for_udim.""" + file_locations_list = [ + FileLocation( + file_path=Path("/path/to/file.1001.exr"), + file_size=1024, + file_hash=None, + ), + FileLocation( + file_path=Path("/path/to/file.1002.exr"), + file_size=1024, + file_hash=None, + ), + FileLocation( + file_path=Path("/path/to/file.1003.exr"), + file_size=1024, + file_hash=None, + ), + ] + + representation = Representation(name="test_1", traits=[ + FileLocations(file_paths=file_locations_list), + UDIM(udim=[1001, 1002, 1003]), + ]) + + udim_trait = representation.get_trait(UDIM) + assert udim_trait.get_file_location_for_udim( + file_locations=representation.get_trait(FileLocations), + udim=1001 + ) == file_locations_list[0]