From 595a3546f37ca82ca30d24407ff5ddc8da1689ce Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ond=C5=99ej=20Samohel?= Date: Sun, 1 Dec 2024 16:11:23 +0100 Subject: [PATCH] :art: add helpers for getting files --- client/ayon_core/pipeline/traits/content.py | 48 ++++++++++++++++++- client/ayon_core/pipeline/traits/time.py | 13 ++++- .../pipeline/traits/test_content_traits.py | 40 ++++++++++++++++ 3 files changed, 98 insertions(+), 3 deletions(-) diff --git a/client/ayon_core/pipeline/traits/content.py b/client/ayon_core/pipeline/traits/content.py index 08dee58383..2808449590 100644 --- a/client/ayon_core/pipeline/traits/content.py +++ b/client/ayon_core/pipeline/traits/content.py @@ -2,10 +2,11 @@ from __future__ import annotations import contextlib +import re # TC003 is there because Path in TYPECHECKING will fail in tests from pathlib import Path # noqa: TC003 -from typing import ClassVar, Optional +from typing import ClassVar, Generator, Optional from pydantic import Field @@ -109,6 +110,51 @@ class FileLocations(TraitBase): id: ClassVar[str] = "ayon.content.FileLocations.v1" file_paths: list[FileLocation] = Field(..., title="File Path") + def get_files(self) -> Generator[Path, None, None]: + """Get all file paths from the trait. + + This method will return all file paths from the trait. + + Yeilds: + Path: List of file paths. + + """ + for file_location in self.file_paths: + yield file_location.file_path + + def get_file_for_frame( + self, + frame: int, + sequence_trait: Optional[Sequence] = None, + ) -> Optional[FileLocation]: + """Get file location for a frame. + + This method will return the file location for a given frame. If the + frame is not found in the file paths, it will return None. + + Args: + frame (int): Frame to get the file location for. + sequence_trait (Sequence): Sequence trait to get the + frame range specs from. + + Returns: + Optional[FileLocation]: File location for the frame. + + """ + frame_regex = r"\.(?P(?P0*)\d+)\.\D+\d?$" + if sequence_trait and sequence_trait.frame_regex: + frame_regex = sequence_trait.frame_regex + + re.compile(frame_regex) + + for file_path in self.get_files(): + result = re.search(frame_regex, file_path.name) + if result: + frame_index = int(result.group("frame")) + if frame_index == frame: + return file_path + return None + def validate(self, representation: Representation) -> None: """Validate the trait. diff --git a/client/ayon_core/pipeline/traits/time.py b/client/ayon_core/pipeline/traits/time.py index 0ab7cfaa9e..22a5c16c13 100644 --- a/client/ayon_core/pipeline/traits/time.py +++ b/client/ayon_core/pipeline/traits/time.py @@ -6,7 +6,7 @@ from enum import Enum, auto from typing import TYPE_CHECKING, ClassVar, Optional import clique -from pydantic import Field +from pydantic import Field, field_validator from .trait import MissingTraitError, TraitBase, TraitValidationError @@ -118,7 +118,7 @@ class Sequence(TraitBase): sequence. frame_padding (int): Frame padding. frame_regex (str): Frame regex - regular expression to match - frame numbers. + frame numbers. Must include 'frame' named group. frame_spec (str): Frame list specification of frames. This takes string like "1-10,20-30,40-50" etc. @@ -132,6 +132,15 @@ class Sequence(TraitBase): frame_regex: Optional[str] = Field(None, title="Frame Regex") frame_spec: Optional[str] = Field(None, title="Frame Specification") + @field_validator("frame_regex") + @classmethod + def validate_frame_regex(cls, v: Optional[str]) -> str: + """Validate frame regex.""" + if v is not None and "?P" not in v: + msg = "Frame regex must include 'frame' named group" + raise ValueError(msg) + return v + def validate(self, representation: Representation) -> None: """Validate the trait.""" super().validate(representation) diff --git a/tests/client/ayon_core/pipeline/traits/test_content_traits.py b/tests/client/ayon_core/pipeline/traits/test_content_traits.py index 065c17a7bb..106b119a66 100644 --- a/tests/client/ayon_core/pipeline/traits/test_content_traits.py +++ b/tests/client/ayon_core/pipeline/traits/test_content_traits.py @@ -139,3 +139,43 @@ def test_file_locations_validation() -> None: with pytest.raises(TraitValidationError): representation.validate() + +def test_get_file_from_frame() -> None: + """Test get_file_from_frame method.""" + file_locations_list = [ + FileLocation( + file_path=Path(f"/path/to/file.{frame}.exr"), + file_size=1024, + file_hash=None, + ) + for frame in range(1001, 1051) + ] + + file_locations_trait: FileLocations = FileLocations( + file_paths=file_locations_list) + + assert file_locations_trait.get_file_for_frame(frame=1001) == \ + file_locations_list[0].file_path + assert file_locations_trait.get_file_for_frame(frame=1050) == \ + file_locations_list[-1].file_path + assert file_locations_trait.get_file_for_frame(frame=1100) is None + + # test with custom regex + sequence = Sequence( + frame_padding=4, + frame_regex=r"boo_(?P\d+)\.exr") + file_locations_list = [ + FileLocation( + file_path=Path(f"/path/to/boo_{frame}.exr"), + file_size=1024, + file_hash=None, + ) + for frame in range(1001, 1051) + ] + + file_locations_trait: FileLocations = FileLocations( + file_paths=file_locations_list) + + assert file_locations_trait.get_file_for_frame( + frame=1001, sequence_trait=sequence) == \ + file_locations_list[0].file_path