Merge remote-tracking branch 'origin/feature/911-new-traits-based-integrator' into feature/911-new-traits-based-integrator

This commit is contained in:
Ondrej Samohel 2024-12-03 10:40:08 +01:00
commit 44f550ea3c
No known key found for this signature in database
GPG key ID: 02376E18990A97C6
7 changed files with 435 additions and 49 deletions

View file

@ -23,7 +23,11 @@ from .time import (
SMPTETimecode,
Static,
)
from .trait import MissingTraitError, TraitBase
from .trait import (
MissingTraitError,
TraitBase,
TraitValidationError,
)
from .two_dimensional import (
UDIM,
Deep,
@ -41,6 +45,7 @@ __all__ = [
"Representation",
"TraitBase",
"MissingTraitError",
"TraitValidationError",
# content
"Bundle",

View file

@ -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
@ -81,7 +82,6 @@ class FileLocation(TraitBase):
file_hash (str): File hash.
"""
name: ClassVar[str] = "FileLocation"
description: ClassVar[str] = "FileLocation Trait Model"
id: ClassVar[str] = "ayon.content.FileLocation.v1"
@ -109,6 +109,50 @@ 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_location_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<frame>(?P<padding>0*)\d+)\.\D+\d?$"
if sequence_trait and sequence_trait.frame_regex:
frame_regex = sequence_trait.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:
frame_index = int(result.group("frame"))
if frame_index == frame:
return location
return None
def validate(self, representation: Representation) -> None:
"""Validate the trait.
@ -131,13 +175,12 @@ class FileLocations(TraitBase):
if representation.contains_trait(FrameRanged):
self._validate_frame_range(representation)
if not representation.contains_trait(Sequence) \
and not representation.contains_trait(Bundle) \
and not representation.contains_trait(UDIM):
# we have multiple files, but it is not a sequence or bundle
# we have multiple files, but it is not a sequence
# or UDIM tile set what it it then? If the files are not related
# to each other then this representation is invalid.
msg = (
"Multiple file locations defined, but no Sequence or Bundle "
"Multiple file locations defined, but no Sequence "
"or UDIM trait defined. If the files are not related to "
"each other, the representation is invalid."
)
@ -351,28 +394,22 @@ class Bundle(TraitBase):
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.
a collection of sub-entities that are part of a single
entity. You can easily reconstruct representations from
the bundle.
Example::
Bundle(
items=[
[
Representation(
traits=[
MimeType(mime_type="image/jpeg"),
FileLocation(file_path="/path/to/file.jpg")
]
)
MimeType(mime_type="image/jpeg"),
FileLocation(file_path="/path/to/file.jpg")
],
[
Representation(
traits=[
MimeType(mime_type="image/png"),
FileLocation(file_path="/path/to/file.png")
]
)
MimeType(mime_type="image/png"),
FileLocation(file_path="/path/to/file.png")
]
]
)

View file

@ -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<frame>" 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)

View file

@ -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,56 @@ 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
def get_udim_from_file_location(
self, file_location: FileLocation) -> Optional[int]:
"""Get UDIM from file location.
Args:
file_location (FileLocation): File location.
Returns:
Optional[int]: UDIM value.
"""
pattern = re.compile(self.udim_regex)
result = re.search(pattern, file_location.file_path.name)
if result:
return int(result.group("udim"))
return None

View file

@ -1,7 +1,10 @@
"""Integrate representations with traits."""
from __future__ import annotations
import contextlib
import copy
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, Any, List
import pyblish.api
@ -22,19 +25,42 @@ from ayon_core.pipeline.publish import (
get_publish_template_name,
)
from ayon_core.pipeline.traits import (
UDIM,
Bundle,
ColorManaged,
FileLocation,
FileLocations,
FrameRanged,
MissingTraitError,
Persistent,
PixelBased,
Representation,
Sequence,
TemplatePath,
TraitValidationError,
)
from pipeline.traits import MissingTraitError, PixelBased
from pipeline.traits.content import FileLocations
if TYPE_CHECKING:
import logging
from pathlib import Path
from pipeline import Anatomy
from ayon_core.pipeline import Anatomy
@dataclass
class TransferItem:
"""Represents single transfer item.
Source file path, destination file path, template that was used to
construct the destination path, template data that was used in the
template, size of the file, checksum of the file.
"""
source: Path
destination: Path
size: int
checksum: str
template: str
template_data: dict[str, Any]
def get_instance_families(instance: pyblish.api.Instance) -> List[str]:
@ -105,9 +131,13 @@ class IntegrateTraits(pyblish.api.InstancePlugin):
order = pyblish.api.IntegratorOrder
log: logging.Logger
def process(self, instance: pyblish.api.Instance) -> None:
def process(self, instance: pyblish.api.Instance) -> None: # noqa: C901, PLR0915, PLR0912
"""Integrate representations with traits.
Todo:
Refactor this method to be more readable and maintainable.
Remove corresponding noqa codes.
Args:
instance (pyblish.api.Instance): Instance to process.
@ -141,7 +171,8 @@ class IntegrateTraits(pyblish.api.InstancePlugin):
return
# 3) get anatomy template
template = self.get_template(instance)
anatomy: Anatomy = instance.context.data["anatomy"]
template: str = self.get_publish_template(instance)
# 4) initialize OperationsSession()
op_session = OperationsSession()
@ -155,30 +186,160 @@ class IntegrateTraits(pyblish.api.InstancePlugin):
)
instance.data["versionEntity"] = version_entity
instance_template_data = {}
transfers = []
# handle {originalDirname} requested in the template
instance_template_data: dict[str, str] = {}
transfers: list[TransferItem] = []
"""
# WIP: This is a draft of the logic that should be implemented
# to handle {originalDirname} in the template
if "{originalDirname}" in template:
instance_template_data = {
"originalDirname": self._get_relative_to_root_original_dirname(
instance)
}
"""
# 6.5) prepare template and data to format it
for representation in representations:
# validate representation first
representation.validate()
# validate representation first, this will go through all traits
# and check if they are valid
try:
representation.validate()
except TraitValidationError as e:
msg = f"Representation '{representation.name}' is invalid: {e}"
raise PublishError(msg) from e
template_data = self.get_template_data_from_representation(
representation, instance)
# add instance based template data
template_data.update(instance_template_data)
if "{originalBasename}" in template:
# Remove 'frame' from template data because original frame
# number will be used.
path_template_object = self.get_publish_template_object(
instance)["path"]
# If representation has FileLocations trait (list of files)
# it can be either Sequence, Bundle or UDIM tile set.
# We do not allow unrelated files in the single representation.
if representation.contains_trait(FileLocations):
# handle sequence
# note: we do not support yet frame sequence of multiple UDIM
# tiles in the same representation
if representation.contains_trait(Sequence):
sequence: Sequence = representation.get_trait(Sequence)
# get the padding from the sequence if the padding on the
# template is higher, us the one from the template
dst_padding = representation.get_trait(
Sequence).frame_padding
frames: list[int] = sequence.get_frame_list(
representation.get_trait(FileLocations),
regex=sequence.frame_regex)
template_padding = anatomy.templates_obj.frame_padding
if template_padding > dst_padding:
dst_padding = template_padding
# go through all frames in the sequence
# find their corresponding file locations
# format their template and add them to transfers
for frame in frames:
template_data["frame"] = frame
template_filled = path_template_object.format_strict(
template_data
)
file_loc: FileLocation = representation.get_trait(
FileLocations).get_file_location_for_frame(
frame, sequence)
transfers.append(
TransferItem(
source=file_loc.file_path,
destination=Path(template_filled),
size=file_loc.file_size,
checksum=file_loc.file_hash,
template=template,
template_data=template_data,
)
)
elif representation.contains_trait(UDIM) and \
not representation.contains_trait(Sequence):
# handle UDIM not in sequence
udim: UDIM = representation.get_trait(UDIM)
for file_loc in representation.get_trait(
FileLocations).file_paths:
template_data["udim"] = (
udim.get_udim_from_file_location(file_loc)
)
template_filled = template.format(**template_data)
transfers.append(
TransferItem(
source=file_loc.file_path,
destination=Path(template_filled),
size=file_loc.file_size,
checksum=file_loc.file_hash,
template=template,
template_data=template_data,
)
)
else:
# This should never happen because the representation
# validation should catch this.
msg = (
"Representation contains FileLocations trait, but "
"is not a Sequence or UDIM."
)
raise PublishError(msg)
elif representation.contains_trait(FileLocation):
# single file representation
template_data.pop("frame", None)
# WIP: use trait logic to get original frame range
# check if files listes in FileLocations trait match frames
# in sequence
with contextlib.suppress(MissingTraitError):
udim = representation.get_trait(UDIM)
template_data["udim"] = udim.udim[0]
template_filled = path_template_object.format_strict(
template_data
)
file_loc: FileLocation = representation.get_trait(FileLocation)
transfers.append(
TransferItem(
source=file_loc.file_path,
destination=Path(template_filled),
size=file_loc.file_size,
checksum=file_loc.file_hash,
template=template,
template_data=template_data,
)
)
elif representation.contains_trait(Bundle):
# handle Bundle
# go through all files in the bundle
pass
# add TemplatePath trait to the representation
representation.add_trait(TemplatePath(
template=template,
data=template_data
))
# format destination path for different types of representations
# in Sequence, we need to handle frame numbering, its padding and
# also the case where it is a UDIM sequence. Note that sequence
# can be non-contiguous.
# --------------------------------
# single file representation or list of non-sequential files is
# simple if representation contains FileLocations trait,
# it is a list of files. there is no hard constrain there,
# but those files should be of the same type ideally - described
# by the same traits.
if representation.contains_trait(Sequence):
# handle template for sequence - this is mostly about
# determining template data for the "udim" and for the "frame".
# Assumption is that the Sequence trait already has the correct
# frame range set. We just need to recalculate to include
# the handles.
...
transfers += self.get_transfers_from_representation(
representation, template, template_data)
@ -270,7 +431,7 @@ class IntegrateTraits(pyblish.api.InstancePlugin):
logger=self.log
)
def get_template(self, instance: pyblish.api.Instance) -> str:
def get_publish_template(self, instance: pyblish.api.Instance) -> str:
"""Return anatomy template name to use for integration.
Args:
@ -287,6 +448,24 @@ class IntegrateTraits(pyblish.api.InstancePlugin):
path_template_obj = publish_template["path"]
return path_template_obj.template.replace("\\", "/")
def get_publish_template_object(
self, instance: pyblish.api.Instance) -> object:
"""Return anatomy template object to use for integration.
Note: What is the actual type of the object?
Args:
instance (pyblish.api.Instance): Instance to process.
Returns:
object: Anatomy template object
"""
# Anatomy data is pre-filled by Collectors
template_name = self.get_template_name(instance)
anatomy = instance.context.data["anatomy"]
return anatomy.get_template_item("publish", template_name)
def prepare_product(
self,
instance: pyblish.api.Instance,
@ -586,12 +765,8 @@ class IntegrateTraits(pyblish.api.InstancePlugin):
template_data["resolution_height"] = representation.get_trait(
PixelBased).display_window_height
# get fps from representation traits
# is this the right way? Isn't it going against the
# trait abstraction?
traits = representation.get_traits()
for trait in traits.values():
if hasattr(trait, "frames_per_second"):
template_data["fps"] = trait.fps
template_data["fps"] = representation.get_trait(
FrameRanged).frames_per_second
# Note: handle "output" and "originalBasename"

View file

@ -139,3 +139,44 @@ def test_file_locations_validation() -> None:
with pytest.raises(TraitValidationError):
representation.validate()
def test_get_file_location_from_frame() -> None:
"""Test get_file_location_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_location_for_frame(frame=1001) == \
file_locations_list[0]
assert file_locations_trait.get_file_location_for_frame(frame=1050) == \
file_locations_list[-1]
assert file_locations_trait.get_file_location_for_frame(frame=1100) is None
# test with custom regex
sequence = Sequence(
frame_padding=4,
frame_regex=r"boo_(?P<frame>\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_location_for_frame(
frame=1001, sequence_trait=sequence) == \
file_locations_list[0]

View file

@ -0,0 +1,62 @@
"""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]
def test_get_udim_from_file_location() -> None:
"""Test get_udim_from_file_location."""
file_location_1 = FileLocation(
file_path=Path("/path/to/file.1001.exr"),
file_size=1024,
file_hash=None,
)
file_location_2 = FileLocation(
file_path=Path("/path/to/file.xxxxx.exr"),
file_size=1024,
file_hash=None,
)
assert UDIM(udim=[1001]).get_udim_from_file_location(
file_location_1) == 1001
assert UDIM(udim=[1001]).get_udim_from_file_location(
file_location_2) is None