extensions-builtin/sd-extension-chainner/nodes/properties/outputs/generic_outputs.py

from __future__ import annotations

from typing import Union

import navi
from nodes.base_output import BaseOutput

from ...impl.color.color import Color
from ...utils.format import format_color_with_channels
from ...utils.seed import Seed


class NumberOutput(BaseOutput):
    def __init__(
        self,
        label: str,
        output_type: navi.ExpressionJson = "number",
    ):
        super().__init__(
            navi.intersect("number", output_type),
            label,
            associated_type=Union[int, float],
        )

    def get_broadcast_type(self, value: int | float):
        return navi.literal(value)

    def enforce(self, value) -> int | float:
        assert isinstance(value, (int, float))
        return value


class TextOutput(BaseOutput):
    def __init__(
        self,
        label: str,
        output_type: navi.ExpressionJson = "string",
    ):
        super().__init__(navi.intersect("string", output_type), label)

    def get_broadcast_type(self, value: str):
        return navi.literal(value)

    def enforce(self, value) -> str:
        assert isinstance(value, str)
        return value


def FileNameOutput(label: str = "Name", of_input: int | None = None):
    output_type = (
        "string"
        if of_input is None
        else f"splitFilePath(Input{of_input}.path).basename"
    )

    return TextOutput(label=label, output_type=output_type)


class SeedOutput(BaseOutput):
    def __init__(self, label: str = "Seed"):
        super().__init__(output_type="Seed", label=label, kind="generic")

    def enforce(self, value) -> Seed:
        assert isinstance(value, Seed)
        return value


class ColorOutput(BaseOutput):
    def __init__(
        self,
        label: str = "Color",
        color_type: navi.ExpressionJson = "Color",
        channels: int | None = None,
    ):
        super().__init__(
            output_type=navi.intersect(color_type, navi.Color(channels=channels)),
            label=label,
            kind="generic",
        )

        self.channels = channels

    def enforce(self, value) -> Color:
        assert isinstance(value, Color)

        if self.channels is not None and value.channels != self.channels:
            expected = format_color_with_channels([self.channels])
            actual = format_color_with_channels([value.channels])
            raise ValueError(
                f"The output {self.label} was supposed to return {expected} but actually returned {actual}."
                f" This is a bug in the implementation of the node."
                f" Please report this bug."
            )

        return value