🎨 udims as a list

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<udim>\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<udim>" 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

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]