From 5748c1593fc824210dcf34f01bc379c2fdbaefd3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20Samohel?= <ondrej.samohel@ynput.io>
Date: Mon, 2 Dec 2024 00:14:49 +0100
Subject: [PATCH 1/2] :art: udims as a list

---
 .../pipeline/traits/two_dimensional.py        | 46 +++++++++++++++++--
 .../traits/test_two_dimesional_traits.py      | 43 +++++++++++++++++
 2 files changed, 86 insertions(+), 3 deletions(-)
 create mode 100644 tests/client/ayon_core/pipeline/traits/test_two_dimesional_traits.py

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<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
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]

From 0d98ff479dca0beb4495daf9594ba6a9637d7c9b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20Samohel?= <ondrej.samohel@ynput.io>
Date: Mon, 2 Dec 2024 00:15:01 +0100
Subject: [PATCH 2/2] :bug: compile regex

---
 client/ayon_core/pipeline/traits/content.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/client/ayon_core/pipeline/traits/content.py b/client/ayon_core/pipeline/traits/content.py
index 2652a8e50c..0ca0cbb3e1 100644
--- a/client/ayon_core/pipeline/traits/content.py
+++ b/client/ayon_core/pipeline/traits/content.py
@@ -144,7 +144,7 @@ class FileLocations(TraitBase):
         if sequence_trait and sequence_trait.frame_regex:
             frame_regex = sequence_trait.frame_regex
 
-        re.compile(frame_regex)
+        frame_regex = re.compile(frame_regex)
         for location in self.file_paths:
             result = re.search(frame_regex, location.file_path.name)
             if result: