D2VWitch ¶

@classmethod
def ensure_obj(
    cls, indexer: str | type[Self] | Self | None = None, /, func_except: FuncExcept | None = None
) -> Self:
    """
    Ensure that the input is a indexer instance, resolving it if necessary.

    Args:
        indexer: Indexer identifier (string, class, or instance). Plugin namespace is also supported.
        func_except: Function returned for custom error handling.

    Returns:
        Indexer instance.
    """
    return _base_ensure_obj(cls, indexer, func_except)

def file_corrupted(self, index_path: SPath) -> None:
    if self.force:
        try:
            index_path.unlink()
        except OSError:
            raise CustomRuntimeError("Index file corrupted, tried to delete it and failed.", self.__class__)
    else:
        raise CustomRuntimeError("Index file corrupted! Delete it and retry.", self.__class__)

@classmethod
def from_param(
    cls, indexer: str | type[Self] | Self | None = None, /, func_except: FuncExcept | None = None
) -> type[Self]:
    """
    Resolve and return an Indexer type from a given input (string, type, or instance).

    Args:
        indexer: Indexer identifier (string, class, or instance). Plugin namespace is also supported.
        func_except: Function returned for custom error handling.

    Returns:
        Resolved indexer type.
    """
    return _base_from_param(cls, indexer, func_except)

def get_cmd(self, files: list[SPath], output: SPath) -> list[str]:
    return list(map(str, [self._get_bin_path(), *files, "--output", output]))

def get_idx_file_path(self, path: SPath) -> SPath:
    return path.with_suffix(f".{self.ext}")

def get_info(self, index_path: SPath, file_idx: int = -1) -> D2VIndexFileInfo:
    with open(index_path) as f:
        file_content = f.read()

    lines = file_content.split("\n")

    head, lines = lines[:2], lines[2:]

    if "DGIndex" not in head[0]:
        self.file_corrupted(index_path)

    raw_header, lines = self._split_lines(self._split_lines(lines)[1])

    header = D2VIndexHeader()

    for rlin in raw_header:
        if split_val := rlin.rstrip().split("="):
            key: str = split_val[0].upper()
            values: list[str] = ",".join(split_val[1:]).split(",")
        else:
            continue

        if key == "STREAM_TYPE":
            header.stream_type = int(values[0])
        elif key == "MPEG_TYPE":
            header.MPEG_type = int(values[0])
        elif key == "IDCT_ALGORITHM":
            header.iDCT_algorithm = int(values[0])
        elif key == "YUVRGB_SCALE":
            header.YUVRGB_scale = int(values[0])
        elif key == "LUMINANCE_FILTER":
            header.luminance_filter = tuple(map(int, values))
        elif key == "CLIPPING":
            header.clipping = list(map(int, values))
        elif key == "ASPECT_RATIO":
            header.aspect = Fraction(*list(map(int, values[0].split(":"))))
        elif key == "PICTURE_SIZE":
            header.pic_size = str(values[0])
        elif key == "FIELD_OPERATION":
            header.field_op = int(values[0])
        elif key == "FRAME_RATE":
            if matches := re.search(r".*\((\d+\/\d+)", values[0]):
                header.frame_rate = Fraction(matches.group(1))
        elif key == "LOCATION":
            header.location = list(map(partial(int, base=16), values))

    frame_data = list[D2VIndexFrameData]()

    if file_idx >= 0:
        for rawline in lines:
            if len(rawline) == 0:
                break

            line = rawline.split(" ", maxsplit=7)

            ffile_idx = int(line[2])

            if ffile_idx < file_idx:
                continue
            elif ffile_idx > file_idx:
                break

            frame_data.append(
                D2VIndexFrameData(
                    int(line[1]),
                    "I",
                    int(line[5]),
                    int(line[6]),
                    int(line[0], 16),
                    int(line[4]),
                    int(line[3]),
                    [int(a, 16) for a in line[7:]],
                )
            )
    elif file_idx == -1:
        for rawline in lines:
            if len(rawline) == 0:
                break

            line = rawline.split(" ")

            frame_data.append(
                D2VIndexFrameData(
                    int(line[1]),
                    "I",
                    int(line[5]),
                    int(line[6]),
                    int(line[0], 16),
                    int(line[4]),
                    int(line[3]),
                    [int(a, 16) for a in line[7:]],
                )
            )

    return D2VIndexFileInfo(index_path, file_idx, header, frame_data)

@classmethod
def get_joined_names(cls, files: list[SPath]) -> str:
    return "_".join([file.name for file in files])

def get_out_folder(
    self, output_folder: SPathLike | Literal[False] | None = None, file: SPath | None = None
) -> SPath:
    if output_folder is None:
        return SPath(file).get_folder() if file else self.get_out_folder(False)

    if not output_folder:
        from tempfile import gettempdir

        return SPath(gettempdir())

    return SPath(output_folder)

def get_video_idx_path(self, folder: SPath, file_hash: str, video_name: SPathLike) -> SPath:
    vid_name = SPath(video_name).stem
    current_indxer = SPath(self._bin_path).name
    filename = "_".join([file_hash, vid_name, current_indxer])

    return self.get_idx_file_path(PackageStorage(folder).get_file(filename))

@classmethod
def get_videos_hash(cls, files: list[SPath]) -> str:
    from hashlib import md5

    length = sum(file.stat().st_size for file in files)
    to_hash = length.to_bytes(32, "little") + cls.get_joined_names(files).encode()
    return md5(to_hash).hexdigest()

def index(
    self,
    files: Sequence[SPath],
    force: bool = False,
    split_files: bool = False,
    output_folder: SPathLike | Literal[False] | None = None,
    *cmd_args: str,
) -> list[SPath]:
    if len(unique_folders := list({f.get_folder().to_str() for f in files})) > 1:
        return [
            c
            for s in (
                self.index(
                    [f for f in files if f.get_folder().to_str() == folder], force, split_files, output_folder
                )
                for folder in unique_folders
            )
            for c in s
        ]

    dest_folder = self.get_out_folder(output_folder, files[0])

    files = sorted(set(files))

    hash_str = self.get_videos_hash(files)

    def _index(files: list[SPath], output: SPath) -> None:
        if output.is_file():
            if output.stat().st_size == 0 or force:
                output.unlink()
            else:
                return self.update_video_filenames(output, files)
        return self._run_index(files, output, cmd_args)

    if not split_files:
        output = self.get_video_idx_path(dest_folder, hash_str, "JOINED" if len(files) > 1 else "SINGLE")
        _index(files, output)
        return [output]

    outputs = [self.get_video_idx_path(dest_folder, hash_str, file.name) for file in files]

    for file, output in zip(files, outputs):
        _index([file], output)

    return outputs

@classmethod
def normalize_filenames(cls, file: SPathLike | Iterable[SPathLike]) -> list[SPath]:
    files = list[SPath]()

    for f in to_arr(file):
        if str(f).startswith("file:///"):
            f = str(f)[8::]

        files.append(SPath(f))

    return files

@inject_self
def source(
    self,
    file: SPathLike | Iterable[SPathLike],
    bits: int | None = None,
    *,
    matrix: MatrixLike | None = None,
    transfer: TransferLike | None = None,
    primaries: PrimariesLike | None = None,
    chroma_location: ChromaLocation | None = None,
    color_range: ColorRangeLike | None = None,
    field_based: FieldBasedLike | None = None,
    idx_props: bool = True,
    **kwargs: Any,
) -> vs.VideoNode:
    index_files = self.index(self.normalize_filenames(file))

    return super().source(
        index_files,
        bits,
        matrix=matrix,
        transfer=transfer,
        primaries=primaries,
        chroma_location=chroma_location,
        color_range=color_range,
        field_based=field_based,
        idx_props=idx_props,
        **kwargs,
    )

@classmethod
def source_func(cls, path: SPathLike, **kwargs: Any) -> vs.VideoNode:
    log.debug("%s: indexing %r; arguments: %r", cls, path, kwargs)
    return cls._source_func(str(path), **kwargs)

def update_video_filenames(self, index_path: SPath, filepaths: list[SPath]) -> None:
    with open(index_path) as file:
        file_content = file.read()

    lines = file_content.split("\n")

    str_filepaths = list(map(str, filepaths))

    if "DGIndex" not in lines[0]:
        self.file_corrupted(index_path)

    if not (n_files := int(lines[1])) or n_files != len(str_filepaths):
        self.file_corrupted(index_path)

    end_videos = lines.index("")

    if lines[2:end_videos] == str_filepaths:
        return

    lines[2:end_videos] = str_filepaths

    with open(index_path, "w") as file:
        file.write("\n".join(lines))