wtd-libreflow-addon/checks/layout/shot.py

import math
import re

import bpy
from mathutils import Matrix

from ... import config, providers
from ...core import Check, Issue, SkipCheck, contexts, is_local

_SUFFIX_RE = re.compile(r"\.\d{3,}$")


class FramerateCheck(Check):
    id = "framerate"
    label = "Frame rate"
    description = "Scene frame rate matches the production frame rate."
    category = "Shot"
    contexts = frozenset({contexts.LAYOUT})
    order = 10
    fix_description = "Set the production frame rate"

    def run(self, ctx):
        render = ctx.scene.render
        fps = render.fps / render.fps_base
        expected = config.FPS / config.FPS_BASE
        if not math.isclose(fps, expected, rel_tol=1e-6):
            yield Issue(f"Frame rate is {fps:g} fps, expected {expected:g}")

    def fix(self, ctx, issues):
        ctx.scene.render.fps = config.FPS
        ctx.scene.render.fps_base = config.FPS_BASE


class ResolutionCheck(Check):
    id = "resolution"
    label = "Resolution"
    description = "Render resolution, percentage and pixel aspect match the production format."
    category = "Shot"
    contexts = frozenset({contexts.LAYOUT})
    order = 20
    fix_description = "Set the production resolution"

    def run(self, ctx):
        render = ctx.scene.render
        size = (render.resolution_x, render.resolution_y)
        if size != tuple(config.RESOLUTION):
            yield Issue(f"Resolution is {size[0]}x{size[1]}, expected {config.RESOLUTION[0]}x{config.RESOLUTION[1]}")
        if render.resolution_percentage != config.RESOLUTION_PERCENTAGE:
            yield Issue(f"Resolution scale is {render.resolution_percentage}%, expected {config.RESOLUTION_PERCENTAGE}%")
        aspect = (render.pixel_aspect_x, render.pixel_aspect_y)
        if any(not math.isclose(a, b) for a, b in zip(aspect, config.PIXEL_ASPECT)):
            yield Issue(f"Pixel aspect is {aspect[0]:g}:{aspect[1]:g}, expected {config.PIXEL_ASPECT[0]:g}:{config.PIXEL_ASPECT[1]:g}")

    def fix(self, ctx, issues):
        render = ctx.scene.render
        render.resolution_x, render.resolution_y = config.RESOLUTION
        render.resolution_percentage = config.RESOLUTION_PERCENTAGE
        render.pixel_aspect_x, render.pixel_aspect_y = config.PIXEL_ASPECT


class FrameRangeCheck(Check):
    id = "frame_range"
    label = "Shot start / end"
    description = "Scene frame range matches the shot range from production tracking (Libreflow / Kitsu)."
    category = "Shot"
    contexts = frozenset({contexts.LAYOUT})
    order = 30
    fix_description = "Set the scene frame range from production tracking"

    def run(self, ctx):
        expected = providers.get_shot_frame_range(ctx)
        if expected is None:
            raise SkipCheck("No shot frame range available: providers.get_shot_frame_range is not connected yet")
        start, end = expected
        scene = ctx.scene
        if (scene.frame_start, scene.frame_end) != (start, end):
            yield Issue(
                f"Frame range is {scene.frame_start}-{scene.frame_end}, expected {start}-{end}",
                data={"range": (start, end)},
            )

    def fix(self, ctx, issues):
        for issue in issues:
            ctx.scene.frame_start, ctx.scene.frame_end = issue.data["range"]


class DuplicateNamesCheck(Check):
    id = "duplicate_names"
    label = "Duplicated names"
    description = "No local datablock carries a '.001' style duplicate suffix."
    category = "Shot"
    contexts = frozenset({contexts.LAYOUT})
    order = 40
    editable_names = True
    fix_description = "Remove the suffix when the base name is free"

    def run(self, ctx):
        for collection_name in config.DUPLICATE_NAME_TYPES:
            collection = getattr(bpy.data, collection_name, None)
            if collection is None:
                continue
            for id_block in collection:
                if is_local(id_block) and _SUFFIX_RE.search(id_block.name):
                    base = _SUFFIX_RE.sub("", id_block.name)
                    clash = " (base name taken)" if base in collection else ""
                    yield Issue.for_id(f"{id_block.name}{clash}", id_block, data={"base": base})

    def fix(self, ctx, issues):
        skipped = []
        for issue in issues:
            id_block = issue.get_id()
            if id_block is None:
                continue
            base = issue.data["base"]
            if base in getattr(bpy.data, issue.id_type):
                skipped.append(f"'{id_block.name}': '{base}' already exists")
                continue
            id_block.name = base
        return skipped


class SetPositionCheck(Check):
    id = "set_positions"
    label = "Set positions"
    description = "Instanced set collections (config SET_COLLECTION_PATTERN) sit at the world origin."
    category = "Shot"
    contexts = frozenset({contexts.LAYOUT})
    order = 50
    fix_description = "Move the set instances back to the origin"

    def run(self, ctx):
        pattern = re.compile(config.SET_COLLECTION_PATTERN)
        found = False
        for ob in ctx.scene.objects:
            collection = ob.instance_collection if ob.instance_type == "COLLECTION" else None
            if collection is None or not pattern.search(collection.name):
                continue
            found = True
            if not _is_identity(ob.matrix_world):
                loc = ob.matrix_world.translation
                yield Issue.for_id(
                    f"{ob.name} ({collection.name}) is not at the origin "
                    f"(location {loc.x:.3f}, {loc.y:.3f}, {loc.z:.3f})",
                    ob,
                )
        if not found:
            raise SkipCheck("No set instance found in the scene")

    def fix(self, ctx, issues):
        for issue in issues:
            ob = issue.get_id()
            if ob is not None:
                ob.matrix_world = Matrix.Identity(4)


def _is_identity(matrix, tolerance=1e-5):
    identity = Matrix.Identity(4)
    return all(abs(matrix[i][j] - identity[i][j]) <= tolerance for i in range(4) for j in range(4))