utils ¶

def depth(
    clip: vs.VideoNode,
    bitdepth: VideoFormatLike | HoldsVideoFormat | int | None = None,
    /,
    sample_type: int | vs.SampleType | None = None,
    *,
    range_in: ColorRangeLike | None = None,
    range_out: ColorRangeLike | None = None,
    dither_type: str | DitherType = DitherType.AUTO,
) -> vs.VideoNode:
    """
    A convenience bitdepth conversion function using only internal plugins if possible.

    This uses exclusively internal plugins except for specific dither_types.
    To check whether your DitherType uses fmtc, use `DitherType.is_fmtc`.

        >>> src_8 = vs.core.std.BlankClip(format=vs.YUV420P8)
        >>> src_10 = depth(src_8, 10)
        >>> src_10.format.name
        'YUV420P10'

        >>> src2_10 = vs.core.std.BlankClip(format=vs.RGB30)
        >>> src2_8 = depth(src2_10, 8, dither_type=Dither.RANDOM)  # override default dither behavior
        >>> src2_8.format.name
        'RGB24'

    Args:
        clip: Input clip.
        bitdepth: Desired bitdepth of the output clip.
        sample_type: Desired sample type of output clip. Allows overriding default float/integer behavior. Accepts
            ``vapoursynth.SampleType`` enums ``vapoursynth.INTEGER`` and ``vapoursynth.FLOAT`` or their values, ``0``
            and ``1`` respectively.
        range_in: Input pixel range (defaults to input `clip`'s range).
        range_out: Output pixel range (defaults to input `clip`'s range).
        dither_type: Dithering algorithm. Allows overriding default dithering behavior.
            See [Dither][vstools.DitherType].

            When integer output is desired but the conversion may produce fractional values,
            defaults to DitherType.VOID if it is available via the fmtc VapourSynth plugin,
            or to Floyd-Steinberg DitherType.ERROR_DIFFUSION for 8-bit output
            or DitherType.ORDERED for higher bit depths.
            In other cases, defaults to DitherType.NONE, or round to nearest.
            See [DitherType.should_dither][vstools.DitherType.should_dither] for more information.

    Returns:
        Converted clip with desired bit depth and sample type. ``ColorFamily`` will be same as input.
    """
    from ..utils import get_video_format

    in_fmt = get_video_format(clip)
    out_fmt = get_video_format(bitdepth or clip, sample_type=sample_type)

    range_out = ColorRange.from_param_or_video(range_out, clip)
    range_in = ColorRange.from_param_or_video(range_in, clip)

    if (in_fmt.bits_per_sample, in_fmt.sample_type, range_in) == (
        out_fmt.bits_per_sample,
        out_fmt.sample_type,
        range_out,
    ):
        return clip

    dither_type = DitherType(dither_type)

    if dither_type is DitherType.AUTO:
        should_dither = DitherType.should_dither(in_fmt, out_fmt, range_in, range_out)

        if hasattr(clip, "fmtc"):
            dither_type = DitherType.VOID
        else:
            dither_type = DitherType.ERROR_DIFFUSION if out_fmt.bits_per_sample == 8 else DitherType.ORDERED
        dither_type = dither_type if should_dither else DitherType.NONE

    new_format = in_fmt.replace(bits_per_sample=out_fmt.bits_per_sample, sample_type=out_fmt.sample_type)

    return dither_type.apply(clip, new_format, range_in, range_out)

def expect_bits(clip: vs.VideoNode, /, expected_depth: int = 16, **kwargs: Any) -> tuple[vs.VideoNode, int]:
    """
    Expected output bitdepth for a clip.

    This function is meant to be used when a clip may not match the expected input bitdepth.
    Both the dithered clip and the original bitdepth are returned.

    Args:
        clip: Input clip.
        expected_depth: Expected bitdepth. Default: 16.

    Returns:
        Tuple containing the clip dithered to the expected depth and the original bitdepth.
    """
    bits = get_depth(clip)

    if bits != expected_depth:
        clip = depth(clip, expected_depth, **kwargs)

    return clip, bits

def flatten_vnodes(*clips: VideoNodeIterable, split_planes: bool = False) -> Sequence[vs.VideoNode]:
    """
    Flatten an array of VideoNodes.

    Args:
        *clips: An array of clips to flatten into a list.
        split_planes: Optionally split the VideoNodes into their individual planes as well. Default: False.

    Returns:
        Flattened list of VideoNodes.
    """
    nodes = list[vs.VideoNode](flatten(clips))

    if not split_planes:
        return nodes

    return reduce(operator.iadd, map(split, nodes), [])

def frame2clip(frame: vs.VideoFrame) -> vs.VideoNode:
    """
    Convert a VideoFrame to a VideoNode.

    Args:
        frame: Input frame.

    Returns:
        1-frame long VideoNode of the input frame.
    """

    key = hash((frame.width, frame.height, frame.format.id))

    if _f2c_cache.get(key, None) is None:
        _f2c_cache[key] = blank_clip = vs.core.std.BlankClip(
            None, frame.width, frame.height, frame.format.id, 1, 1, 1, [0] * frame.format.num_planes, True
        )
    else:
        blank_clip = _f2c_cache[key]

    frame_cp = frame.copy()

    return vs.core.std.ModifyFrame(blank_clip, blank_clip, lambda n, f: frame_cp)

def get_b(clip: vs.VideoNode, /) -> vs.VideoNode:
    """
    Extract the blue plane of the given clip.

    Args:
        clip: Input clip.

    Returns:
        B plane of the input clip.

    Raises:
        CustomValueError: Clip is not RGB.
    """

    InvalidColorFamilyError.check(clip, vs.RGB, get_b)

    return plane(clip, 2)

def get_g(clip: vs.VideoNode, /) -> vs.VideoNode:
    """
    Extract the green plane of the given clip.

    Args:
        clip: Input clip.

    Returns:
        G plane of the input clip.

    Raises:
        CustomValueError: Clip is not RGB.
    """

    InvalidColorFamilyError.check(clip, vs.RGB, get_g)

    return plane(clip, 1)

def get_r(clip: vs.VideoNode, /) -> vs.VideoNode:
    """
    Extract the red plane of the given clip.

    Args:
        clip: Input clip.

    Returns:
        R plane of the input clip.

    Raises:
        CustomValueError: Clip is not RGB.
    """

    InvalidColorFamilyError.check(clip, vs.RGB, get_r)

    return plane(clip, 0)

def get_u(clip: vs.VideoNode, /) -> vs.VideoNode:
    """
    Extract the first chroma (U) plane of the given clip.

    Args:
        clip: Input clip.

    Returns:
        Y plane of the input clip.

    Raises:
        CustomValueError: Clip is not YUV.
    """

    InvalidColorFamilyError.check(clip, vs.YUV, get_u)

    return plane(clip, 1)

def get_v(clip: vs.VideoNode, /) -> vs.VideoNode:
    """
    Extract the second chroma (V) plane of the given clip.

    Args:
        clip: Input clip.

    Returns:
        V plane of the input clip.

    Raises:
        CustomValueError: Clip is not YUV.
    """

    InvalidColorFamilyError.check(clip, vs.YUV, get_v)

    return plane(clip, 2)

def get_y(clip: vs.VideoNode, /) -> vs.VideoNode:
    """
    Extract the luma (Y) plane of the given clip.

    Args:
        clip: Input clip.

    Returns:
        Y plane of the input clip.

    Raises:
        CustomValueError: Clip is not GRAY or YUV.
    """

    InvalidColorFamilyError.check(clip, [vs.YUV, vs.GRAY], get_y)

    return plane(clip, 0)

def insert_clip(clip: vs.VideoNode, /, insert: vs.VideoNode, start_frame: int, strict: bool = True) -> vs.VideoNode:
    """
    Replace frames of a longer clip with those of a shorter one.

    The insert clip may NOT exceed the final frame of the input clip.
    This limitation can be circumvented by setting `strict=False`.

    Args:
        clip: Input clip.
        insert: Clip to insert into the input clip.
        start_frame: Frame to start inserting from.
        strict: Throw an error if the inserted clip exceeds the final frame of the input clip. If False, truncate the
            inserted clip instead. Default: True.

    Returns:
        Clip with frames replaced by the insert clip.

    Raises:
        CustomValueError: Insert clip is too long, ``strict=False``, and exceeds the final frame of the input clip.
    """

    if start_frame == 0:
        return vs.core.std.Splice([insert, clip[insert.num_frames :]])

    pre = clip[:start_frame]
    insert_diff = (start_frame + insert.num_frames) - clip.num_frames

    if insert_diff == 0:
        return vs.core.std.Splice([pre, insert])

    if insert_diff < 0:
        return vs.core.std.Splice([pre, insert, clip[insert_diff:]])

    if strict:
        raise ClipLengthError(
            "Inserted clip is too long and exceeds the final frame of the input clip.",
            insert_clip,
            {"clip": clip.num_frames, "diff": insert_diff},
        )

    return vs.core.std.Splice([pre, insert[:-insert_diff]])

def join(
    clips: vs.VideoNode | Sequence[vs.VideoNode] | Mapping[Planes, vs.VideoNode | None],
    plane1_or_family: vs.VideoNode | vs.ColorFamily | None = None,
    plane2: vs.VideoNode | None = None,
    alpha: vs.VideoNode | None = None,
    /,
    *,
    family: vs.ColorFamily = vs.ColorFamily.YUV,
    prop_src: vs.VideoNode | SupportsIndex | None = None,
) -> vs.VideoNode:
    """
    General interface to combine clips or planes into a single clip.

    Args:
        clips: First plane, sequence of single-plane clips or mapping of planes to clips:
        plane1_or_family: Chroma clip, second plane or color family, depending on usage.
        plane2: Third plane when combining three planes.
        alpha: Optional alpha clip.
        family: Output color family. Defaults to YUV.
        prop_src: Optional clip or index to copy frame properties from.

    Raises:
        CustomIndexError: Invalid plane index.
        CustomTypeError: Invalid input type.

    Returns:
        Clip with combined planes.
    """
    if isinstance(clips, Mapping):
        from ..functions import flatten_vnodes

        clips_map = dict[int, vs.VideoNode]()

        for p_key, node in clips.items():
            if node is None:
                continue

            if p_key is None:
                clips_map.update(enumerate(flatten_vnodes(node, split_planes=True)))
            else:
                clips_map.update((i, plane(node, i)) for i in to_arr(p_key))

        return join([clips_map[i] for i in sorted(clips_map)], family, prop_src=prop_src)

    if not isinstance(clips, vs.VideoNode):
        if isinstance(plane1_or_family, vs.ColorFamily):
            family = plane1_or_family

        if (n_clips := len(clips)) == 1:
            return clips[0]
        if n_clips in (2, 3, 4):
            return join(*clips, family=family, prop_src=prop_src)
        else:
            raise CustomIndexError("Too many or not enough clips/planes passed!", join, n_clips)

    if not isinstance(plane1_or_family, vs.VideoNode):
        raise CustomTypeError(func=join)

    if not plane2:
        InvalidColorFamilyError.check(
            (family, plane1_or_family),
            vs.YUV,
            join,
            "When combining two clips, color family and chroma clip must be {correct}, not {wrong}.",
        )

        clips = [clips, plane1_or_family]
        planes = [0, 1, 2]

    else:
        clips = [clips, plane1_or_family, plane2]
        planes = [0, 0, 0]

    if not isinstance(prop_src, (vs.VideoNode, NoneType)):
        prop_src = clips[prop_src.__index__()]

    joined = vs.core.std.ShufflePlanes(clips, planes, family, prop_src)

    if alpha:
        joined = joined.std.ClipToProp(alpha, "_Alpha")

    return joined

def limiter[**P](
    clip_or_func: vs.VideoNode | Callable[P, vs.VideoNode] | None = None,
    /,
    min_val: float | Sequence[float] | None = None,
    max_val: float | Sequence[float] | None = None,
    *,
    tv_range: bool = False,
    mask: bool = False,
    planes: Planes = None,
    func: FuncExcept | None = None,
) -> Union[
    vs.VideoNode,
    Callable[P, vs.VideoNode],
    Callable[[Callable[P, vs.VideoNode]], Callable[P, vs.VideoNode]],
]:
    """
    Wraps `vs-zip <https://github.com/dnjulek/vapoursynth-zip>`.Limiter but only processes
    if clip format is not integer, a min/max val is specified or tv_range is True.

    Args:
        clip_or_func: Clip to process or function that returns a VideoNode to be processed.
        min_val: Lower bound. Defaults to the lowest allowed value for the input. Can be specified for each plane
            individually.
        max_val: Upper bound. Defaults to the highest allowed value for the input. Can be specified for each plane
            individually.
        tv_range: Changes min/max defaults values to LIMITED.
        mask: Float chroma range from -0.5/0.5 to 0.0/1.0.
        planes: Planes to process.
        func: Function returned for custom error handling. This should only be set by VS package developers.

    Returns:
        Clamped clip.
    """
    if callable(clip_or_func):
        _func = clip_or_func

        @wraps(_func)
        def _wrapper(*args: P.args, **kwargs: P.kwargs) -> vs.VideoNode:
            return limiter(
                _func(*args, **kwargs),
                min_val,
                max_val,
                tv_range=tv_range,
                mask=mask,
                planes=planes,
                func=func or _func,
            )

        return _wrapper

    func = func or limiter
    clip = clip_or_func

    if clip is None:
        return partial(limiter, min_val=min_val, max_val=max_val, tv_range=tv_range, planes=planes, func=func)

    if all([clip.format.sample_type == vs.INTEGER, min_val is None, max_val is None, tv_range is False]):
        return clip

    if not (min_val == max_val is None):
        from ..utils import get_lowest_values, get_peak_values

        min_val = normalize_seq(min_val or get_lowest_values(clip, clip), clip.format.num_planes)
        max_val = normalize_seq(max_val or get_peak_values(clip, clip), clip.format.num_planes)

    return clip.vszip.Limiter(min_val, max_val, tv_range, mask, planes)

def plane(clip: vs.VideoNode, index: SupportsIndex, /, strict: bool = True) -> vs.VideoNode:
    """
    Extract a plane from the given clip.

    Args:
        clip: Input clip.
        index: Index of the plane to extract.
        strict: If False, removes `_Matrix` property when the input clip is RGB.

    Returns:
        Grayscale clip of the clip's plane.
    """
    if clip.format.num_planes == 1 and index.__index__() == 0:
        return clip

    if not strict and clip.format.color_family is vs.RGB:
        clip = vs.core.std.RemoveFrameProps(clip, "_Matrix")

    return vs.core.std.ShufflePlanes(clip, index.__index__(), vs.GRAY)

def split(clip: vs.VideoNode, /, strict: bool = True) -> list[vs.VideoNode]:
    """
    Split a clip into a list of individual planes.

    Args:
        clip: Input clip.
        strict: If False, removes `_Matrix` property when the input clip is RGB.

    Returns:
        List of individual planes.
    """
    return [clip] if clip.format.num_planes == 1 else [plane(clip, i, strict) for i in range(clip.format.num_planes)]

def stack_clips(clips: Iterable[VideoNodeIterable]) -> vs.VideoNode:
    """
    Recursively stack clips in alternating directions: horizontal → vertical → horizontal → ...

    This function takes a nested sequence of clips (or lists of clips) and stacks them
    alternately along the horizontal and vertical axes at each nesting level.

    Examples:
        - Stack a list of clips horizontally:
          ```py
          from vstools import core, split, vs

          clip = core.std.BlankClip(format=vs.RGB24)
          clips = split(clip)

          stacked = stack_clips(clips)
          ```

        - Stack a list of clips vertically (wrap in another list):
          ```py
          from vstools import core, split, vs

          clip = core.std.BlankClip(format=vs.RGB24)
          clips = split(clip)

          stacked = stack_clips([clips])
          ```

        - Stack the Y plane horizontally with the U and V planes stacked vertically:
          ```py
          from vstools import core, split, vs

          clip = core.std.BlankClip(format=vs.YUV420P8)
          y, u, v = split(clip)

          stacked = stack_clips([y, [u, v]])
          ```

        - Stack multiple YUV clips, with Y planes horizontally and UV planes vertically:
          ```py
          from vstools import core, split, vs

          yuv_clips = [...]  # all must share format and height

          clips = []
          for yuv_clip in yuv_clips:
              y, *uv = split(yuv_clip)
              clips.extend([y, uv])

          stacked = stack_clips(clips)
          ```

        - Using ``append`` instead of ``extend`` (and wrapping the sequence, e.g. ``stack_clips([clips])``)
          changes the stacking layout, since it alters the nesting depth.

    Args:
        clips: A (possibly nested) sequence of clips to be stacked.

    Returns:
        Stacked clips.
    """
    return vs.core.std.StackHorizontal(
        [
            inner_clips
            if isinstance(inner_clips, vs.VideoNode)
            else (
                vs.core.std.StackVertical(
                    [clipa if isinstance(clipa, vs.VideoNode) else stack_clips(clipa) for clipa in inner_clips]
                )
            )
            for inner_clips in clips
        ]
    )