Use data class for TempData to get some type hints

This commit is contained in:
Roy Nieterau 2025-06-11 20:38:03 +02:00
parent d20ade6d79
commit c74fbc2aff

View file

@ -1,3 +1,4 @@
from __future__ import annotations
import os
import re
import copy
@ -6,6 +7,7 @@ import shutil
import subprocess
from abc import ABC, abstractmethod
from typing import Dict, Any, Optional
from dataclasses import dataclass, field
import tempfile
import clique
@ -35,6 +37,39 @@ from ayon_core.pipeline.publish import (
from ayon_core.pipeline.publish.lib import add_repre_files_for_cleanup
@dataclass
class TempData:
"""Temporary data used across extractor's process."""
fps: float
frame_start: int
frame_end: int
handle_start: int
handle_end: int
frame_start_handle: int
frame_end_handle: int
output_frame_start: int
output_frame_end: int
pixel_aspect: float
resolution_width: int
resolution_height: int
origin_repre: dict[str, Any]
input_is_sequence: bool
first_sequence_frame: int
input_allow_bg: bool
with_audio: bool
without_handles: bool
handles_are_set: bool
input_ext: str
explicit_input_paths: list[str]
paths_to_remove: list[str]
# Set later
full_output_path: str = ""
filled_files: Dict[int, str] = field(default_factory=list)
output_ext_is_image: bool = True
output_is_sequence: bool = True
def frame_to_timecode(frame: int, fps: float) -> str:
"""Convert a frame number and FPS to editorial timecode (HH:MM:SS:FF).
@ -405,10 +440,10 @@ class ExtractReview(pyblish.api.InstancePlugin):
temp_data = self.prepare_temp_data(instance, repre, output_def)
new_frame_files = {}
if temp_data["input_is_sequence"]:
if temp_data.input_is_sequence:
self.log.debug("Checking sequence to fill gaps in sequence..")
files = temp_data["origin_repre"]["files"]
files = temp_data.origin_repre["files"]
collections = clique.assemble(
files,
)[0]
@ -423,18 +458,18 @@ class ExtractReview(pyblish.api.InstancePlugin):
new_frame_files = self.fill_sequence_gaps_from_existing(
collection=collection,
staging_dir=new_repre["stagingDir"],
start_frame=temp_data["frame_start"],
end_frame=temp_data["frame_end"],
start_frame=temp_data.frame_start,
end_frame=temp_data.frame_end,
)
elif fill_missing_frames == "blank":
new_frame_files = self.fill_sequence_gaps_with_blanks(
collection=collection,
staging_dir=new_repre["stagingDir"],
start_frame=temp_data["frame_start"],
end_frame=temp_data["frame_end"],
resolution_width=temp_data["resolution_width"],
resolution_height=temp_data["resolution_height"],
extension=temp_data["input_ext"],
start_frame=temp_data.frame_start,
end_frame=temp_data.frame_end,
resolution_width=temp_data.resolution_width,
resolution_height=temp_data.resolution_height,
extension=temp_data.input_ext,
temp_data=temp_data
)
elif fill_missing_frames == "previous_version":
@ -443,8 +478,8 @@ class ExtractReview(pyblish.api.InstancePlugin):
staging_dir=new_repre["stagingDir"],
instance=instance,
current_repre_name=repre["name"],
start_frame=temp_data["frame_start"],
end_frame=temp_data["frame_end"],
start_frame=temp_data.frame_start,
end_frame=temp_data.frame_end,
)
# fallback to original workflow
if new_frame_files is None:
@ -452,11 +487,11 @@ class ExtractReview(pyblish.api.InstancePlugin):
self.fill_sequence_gaps_from_existing(
collection=collection,
staging_dir=new_repre["stagingDir"],
start_frame=temp_data["frame_start"],
end_frame=temp_data["frame_end"],
start_frame=temp_data.frame_start,
end_frame=temp_data.frame_end,
))
elif fill_missing_frames == "only_rendered":
temp_data["explicit_input_paths"] = [
temp_data.explicit_input_paths = [
os.path.join(
new_repre["stagingDir"], file
).replace("\\", "/")
@ -467,10 +502,10 @@ class ExtractReview(pyblish.api.InstancePlugin):
# modify range for burnins
instance.data["frameStart"] = frame_start
instance.data["frameEnd"] = frame_end
temp_data["frame_start"] = frame_start
temp_data["frame_end"] = frame_end
temp_data.frame_start = frame_start
temp_data.frame_end = frame_end
temp_data["filled_files"] = new_frame_files
temp_data.filled_files = new_frame_files
# create or update outputName
output_name = new_repre.get("outputName", "")
@ -478,7 +513,7 @@ class ExtractReview(pyblish.api.InstancePlugin):
if output_name:
output_name += "_"
output_name += output_def["filename_suffix"]
if temp_data["without_handles"]:
if temp_data.without_handles:
output_name += "_noHandles"
# add outputName to anatomy format fill_data
@ -491,7 +526,7 @@ class ExtractReview(pyblish.api.InstancePlugin):
# like Resolve or Premiere can detect the start frame for e.g.
# review output files
"timecode": frame_to_timecode(
frame=temp_data["frame_start_handle"],
frame=temp_data.frame_start_handle,
fps=float(instance.data["fps"])
)
})
@ -508,7 +543,7 @@ class ExtractReview(pyblish.api.InstancePlugin):
except ZeroDivisionError:
# TODO recalculate width and height using OIIO before
# conversion
if 'exr' in temp_data["origin_repre"]["ext"]:
if 'exr' in temp_data.origin_repre["ext"]:
self.log.warning(
(
"Unsupported compression on input files."
@ -531,16 +566,16 @@ class ExtractReview(pyblish.api.InstancePlugin):
for filepath in new_frame_files.values():
os.unlink(filepath)
for filepath in temp_data["paths_to_remove"]:
for filepath in temp_data.paths_to_remove:
os.unlink(filepath)
new_repre.update({
"fps": temp_data["fps"],
"fps": temp_data.fps,
"name": "{}_{}".format(output_name, output_ext),
"outputName": output_name,
"outputDef": output_def,
"frameStartFtrack": temp_data["output_frame_start"],
"frameEndFtrack": temp_data["output_frame_end"],
"frameStartFtrack": temp_data.output_frame_start,
"frameEndFtrack": temp_data.output_frame_end,
"ffmpeg_cmd": subprcs_cmd
})
@ -566,7 +601,7 @@ class ExtractReview(pyblish.api.InstancePlugin):
# - there can be more than one collection
return isinstance(repre["files"], (list, tuple))
def prepare_temp_data(self, instance, repre, output_def):
def prepare_temp_data(self, instance, repre, output_def) -> TempData:
"""Prepare dictionary with values used across extractor's process.
All data are collected from instance, context, origin representation
@ -582,7 +617,7 @@ class ExtractReview(pyblish.api.InstancePlugin):
output_def (dict): Definition of output of this plugin.
Returns:
dict: All data which are used across methods during process.
TempData: All data which are used across methods during process.
Their values should not change during process but new keys
with values may be added.
"""
@ -647,30 +682,30 @@ class ExtractReview(pyblish.api.InstancePlugin):
else:
ext = os.path.splitext(repre["files"])[1].replace(".", "")
return {
"fps": float(instance.data["fps"]),
"frame_start": frame_start,
"frame_end": frame_end,
"handle_start": handle_start,
"handle_end": handle_end,
"frame_start_handle": frame_start_handle,
"frame_end_handle": frame_end_handle,
"output_frame_start": int(output_frame_start),
"output_frame_end": int(output_frame_end),
"pixel_aspect": instance.data.get("pixelAspect", 1),
"resolution_width": instance.data.get("resolutionWidth"),
"resolution_height": instance.data.get("resolutionHeight"),
"origin_repre": repre,
"input_is_sequence": input_is_sequence,
"first_sequence_frame": first_sequence_frame,
"input_allow_bg": input_allow_bg,
"with_audio": with_audio,
"without_handles": without_handles,
"handles_are_set": handles_are_set,
"input_ext": ext,
"explicit_input_paths": [], # absolute paths to rendered files
"paths_to_remove": []
}
return TempData(
fps=float(instance.data["fps"]),
frame_start=frame_start,
frame_end=frame_end,
handle_start=handle_start,
handle_end=handle_end,
frame_start_handle=frame_start_handle,
frame_end_handle=frame_end_handle,
output_frame_start=int(output_frame_start),
output_frame_end=int(output_frame_end),
pixel_aspect=instance.data.get("pixelAspect", 1),
resolution_width=instance.data.get("resolutionWidth"),
resolution_height=instance.data.get("resolutionHeight"),
origin_repre=repre,
input_is_sequence=input_is_sequence,
first_sequence_frame=first_sequence_frame,
input_allow_bg=input_allow_bg,
with_audio=with_audio,
without_handles=without_handles,
handles_are_set=handles_are_set,
input_ext=ext,
explicit_input_paths=[], # absolute paths to rendered files
paths_to_remove=[]
)
def _ffmpeg_arguments(
self,
@ -691,7 +726,7 @@ class ExtractReview(pyblish.api.InstancePlugin):
instance (Instance): Currently processed instance.
new_repre (dict): Representation representing output of this
process.
temp_data (dict): Base data for successful process.
temp_data (TempData): Base data for successful process.
"""
# Get FFmpeg arguments from profile presets
@ -733,32 +768,32 @@ class ExtractReview(pyblish.api.InstancePlugin):
# Set output frames len to 1 when output is single image
if (
temp_data["output_ext_is_image"]
and not temp_data["output_is_sequence"]
temp_data.output_ext_is_image
and not temp_data.output_is_sequence
):
output_frames_len = 1
else:
output_frames_len = (
temp_data["output_frame_end"]
- temp_data["output_frame_start"]
temp_data.output_frame_end
- temp_data.output_frame_start
+ 1
)
duration_seconds = float(output_frames_len / temp_data["fps"])
duration_seconds = float(output_frames_len / temp_data.fps)
# Define which layer should be used
if layer_name:
ffmpeg_input_args.extend(["-layer", layer_name])
explicit_input_paths = temp_data["explicit_input_paths"]
if temp_data["input_is_sequence"] and not explicit_input_paths:
explicit_input_paths = temp_data.explicit_input_paths
if temp_data.input_is_sequence and not explicit_input_paths:
# Set start frame of input sequence (just frame in filename)
# - definition of input filepath
# - add handle start if output should be without handles
start_number = temp_data["first_sequence_frame"]
if temp_data["without_handles"] and temp_data["handles_are_set"]:
start_number += temp_data["handle_start"]
start_number = temp_data.first_sequence_frame
if temp_data.without_handles and temp_data.handles_are_set:
start_number += temp_data.handle_start
ffmpeg_input_args.extend([
"-start_number", str(start_number)
])
@ -771,32 +806,32 @@ class ExtractReview(pyblish.api.InstancePlugin):
# }
# Add framerate to input when input is sequence
ffmpeg_input_args.extend([
"-framerate", str(temp_data["fps"])
"-framerate", str(temp_data.fps)
])
# Add duration of an input sequence if output is video
if not temp_data["output_is_sequence"]:
if not temp_data.output_is_sequence:
ffmpeg_input_args.extend([
"-to", "{:0.10f}".format(duration_seconds)
])
if temp_data["output_is_sequence"] and not explicit_input_paths:
if temp_data.output_is_sequence and not explicit_input_paths:
# Set start frame of output sequence (just frame in filename)
# - this is definition of an output
ffmpeg_output_args.extend([
"-start_number", str(temp_data["output_frame_start"])
"-start_number", str(temp_data.output_frame_start)
])
# Change output's duration and start point if should not contain
# handles
if temp_data["without_handles"] and temp_data["handles_are_set"]:
if temp_data.without_handles and temp_data.handles_are_set:
# Set output duration in seconds
ffmpeg_output_args.extend([
"-t", "{:0.10}".format(duration_seconds)
])
# Add -ss (start offset in seconds) if input is not sequence
if not temp_data["input_is_sequence"]:
start_sec = float(temp_data["handle_start"]) / temp_data["fps"]
if not temp_data.input_is_sequence:
start_sec = float(temp_data.handle_start) / temp_data.fps
# Set start time without handles
# - Skip if start sec is 0.0
if start_sec > 0.0:
@ -805,7 +840,7 @@ class ExtractReview(pyblish.api.InstancePlugin):
])
# Set frame range of output when input or output is sequence
elif temp_data["output_is_sequence"]:
elif temp_data.output_is_sequence:
ffmpeg_output_args.extend([
"-frames:v", str(output_frames_len)
])
@ -813,10 +848,10 @@ class ExtractReview(pyblish.api.InstancePlugin):
if not explicit_input_paths:
# Add video/image input path
ffmpeg_input_args.extend([
"-i", path_to_subprocess_arg(temp_data["full_input_path"])
"-i", path_to_subprocess_arg(temp_data.full_input_path)
])
else:
frame_duration = 1 / temp_data["fps"]
frame_duration = 1 / temp_data.fps
explicit_frames_meta = tempfile.NamedTemporaryFile(
mode="w", prefix="explicit_frames", suffix=".txt", delete=False
@ -826,21 +861,21 @@ class ExtractReview(pyblish.api.InstancePlugin):
with open(explicit_frames_path, "w") as fp:
lines = [
f"file '{path}'{os.linesep}duration {frame_duration}"
for path in temp_data["explicit_input_paths"]
for path in temp_data.explicit_input_paths
]
fp.write("\n".join(lines))
temp_data["paths_to_remove"].append(explicit_frames_path)
temp_data.paths_to_remove.append(explicit_frames_path)
# let ffmpeg use only rendered files, might have gaps
ffmpeg_input_args.extend([
"-f", "concat",
"-safe", "0",
"-i", path_to_subprocess_arg(explicit_frames_path),
"-r", str(temp_data["fps"])
"-r", str(temp_data.fps)
])
# Add audio arguments if there are any. Skipped when output are images.
if not temp_data["output_ext_is_image"] and temp_data["with_audio"]:
if not temp_data.output_ext_is_image and temp_data.with_audio:
audio_in_args, audio_filters, audio_out_args = self.audio_args(
instance, temp_data, duration_seconds
)
@ -862,7 +897,7 @@ class ExtractReview(pyblish.api.InstancePlugin):
bg_red, bg_green, bg_blue, bg_alpha = bg_color
if bg_alpha > 0.0:
if not temp_data["input_allow_bg"]:
if not temp_data.input_allow_bg:
self.log.info((
"Output definition has defined BG color input was"
" resolved as does not support adding BG."
@ -893,7 +928,7 @@ class ExtractReview(pyblish.api.InstancePlugin):
# NOTE This must be latest added item to output arguments.
ffmpeg_output_args.append(
path_to_subprocess_arg(temp_data["full_output_path"])
path_to_subprocess_arg(temp_data.full_output_path)
)
return self.ffmpeg_full_args(
@ -1072,7 +1107,7 @@ class ExtractReview(pyblish.api.InstancePlugin):
resolution_width: int,
resolution_height: int,
extension: str,
temp_data: Dict[str, Any]
temp_data: TempData
) -> Optional[Dict[int, str]]:
"""Fills missing files by blank frame."""
@ -1089,7 +1124,7 @@ class ExtractReview(pyblish.api.InstancePlugin):
blank_frame_path = self._create_blank_frame(
staging_dir, extension, resolution_width, resolution_height
)
temp_data["paths_to_remove"].append(blank_frame_path)
temp_data.paths_to_remove.append(blank_frame_path)
speedcopy.copyfile(blank_frame_path, hole_fpath)
added_files[frame] = hole_fpath
@ -1176,7 +1211,7 @@ class ExtractReview(pyblish.api.InstancePlugin):
return added_files
def input_output_paths(self, new_repre, output_def, temp_data):
def input_output_paths(self, new_repre, output_def, temp_data: TempData):
"""Deduce input nad output file paths based on entered data.
Input may be sequence of images, video file or single image file and
@ -1189,11 +1224,11 @@ class ExtractReview(pyblish.api.InstancePlugin):
"sequence_file" (if output is sequence) keys to new representation.
"""
repre = temp_data["origin_repre"]
repre = temp_data.origin_repre
src_staging_dir = repre["stagingDir"]
dst_staging_dir = new_repre["stagingDir"]
if temp_data["input_is_sequence"]:
if temp_data.input_is_sequence:
collections = clique.assemble(repre["files"])[0]
full_input_path = os.path.join(
src_staging_dir,
@ -1218,13 +1253,13 @@ class ExtractReview(pyblish.api.InstancePlugin):
# Make sure to have full path to one input file
full_input_path_single_file = full_input_path
filled_files = temp_data["filled_files"]
filled_files = temp_data.filled_files
if filled_files:
first_frame, first_file = next(iter(filled_files.items()))
if first_file < full_input_path_single_file:
self.log.warning(f"Using filled frame: '{first_file}'")
full_input_path_single_file = first_file
temp_data["first_sequence_frame"] = first_frame
temp_data.first_sequence_frame = first_frame
filename_suffix = output_def["filename_suffix"]
@ -1252,8 +1287,8 @@ class ExtractReview(pyblish.api.InstancePlugin):
)
if output_is_sequence:
new_repre_files = []
frame_start = temp_data["output_frame_start"]
frame_end = temp_data["output_frame_end"]
frame_start = temp_data.output_frame_start
frame_end = temp_data.output_frame_end
filename_base = "{}_{}".format(filename, filename_suffix)
# Temporary template for frame filling. Example output:
@ -1290,18 +1325,18 @@ class ExtractReview(pyblish.api.InstancePlugin):
new_repre["stagingDir"] = dst_staging_dir
# Store paths to temp data
temp_data["full_input_path"] = full_input_path
temp_data["full_input_path_single_file"] = full_input_path_single_file
temp_data["full_output_path"] = full_output_path
temp_data.full_input_path = full_input_path
temp_data.full_input_path_single_file = full_input_path_single_file
temp_data.full_output_path = full_output_path
# Store information about output
temp_data["output_ext_is_image"] = output_ext_is_image
temp_data["output_is_sequence"] = output_is_sequence
temp_data.output_ext_is_image = output_ext_is_image
temp_data.output_is_sequence = output_is_sequence
self.log.debug("Input path {}".format(full_input_path))
self.log.debug("Output path {}".format(full_output_path))
def audio_args(self, instance, temp_data, duration_seconds):
def audio_args(self, instance, temp_data: TempData, duration_seconds):
"""Prepares FFMpeg arguments for audio inputs."""
audio_in_args = []
audio_filters = []
@ -1318,7 +1353,7 @@ class ExtractReview(pyblish.api.InstancePlugin):
frame_start_ftrack = instance.data.get("frameStartFtrack")
if frame_start_ftrack is not None:
offset_frames = frame_start_ftrack - audio["offset"]
offset_seconds = offset_frames / temp_data["fps"]
offset_seconds = offset_frames / temp_data.fps
if offset_seconds > 0:
audio_in_args.append(
@ -1502,7 +1537,7 @@ class ExtractReview(pyblish.api.InstancePlugin):
return output
def rescaling_filters(self, temp_data, output_def, new_repre):
def rescaling_filters(self, temp_data: TempData, output_def, new_repre):
"""Prepare vieo filters based on tags in new representation.
It is possible to add letterboxes to output video or rescale to
@ -1522,7 +1557,7 @@ class ExtractReview(pyblish.api.InstancePlugin):
self.log.debug("reformat_in_baking: `{}`".format(reformat_in_baking))
# NOTE Skipped using instance's resolution
full_input_path_single_file = temp_data["full_input_path_single_file"]
full_input_path_single_file = temp_data.full_input_path_single_file
try:
streams = get_ffprobe_streams(
full_input_path_single_file, self.log
@ -1547,7 +1582,7 @@ class ExtractReview(pyblish.api.InstancePlugin):
break
# Get instance data
pixel_aspect = temp_data["pixel_aspect"]
pixel_aspect = temp_data.pixel_aspect
if reformat_in_baking:
self.log.debug((
"Using resolution from input. It is already "
@ -1642,8 +1677,8 @@ class ExtractReview(pyblish.api.InstancePlugin):
# - use instance resolution only if there were not scale changes
# that may massivelly affect output 'use_input_res'
if not use_input_res and output_width is None or output_height is None:
output_width = temp_data["resolution_width"]
output_height = temp_data["resolution_height"]
output_width = temp_data.resolution_width
output_height = temp_data.resolution_height
# Use source's input resolution instance does not have set it.
if output_width is None or output_height is None: