ytptube/app/features/ytdlp/ytdlp.py

# flake8: noqa: F401, RUF100, W291, I001
import sys
from typing import Any

import yt_dlp
from yt_dlp.globals import extractors as ytdlp_extractors
from yt_dlp.utils import make_archive_id

from app.features.ytdlp.outtmpl import rewrite_outtmpl
from app.features.ytdlp.patches import apply_ytdlp_patches
from app.library.cf_solver_handler import set_cf_handler
from app.library.log import get_logger


class _ArchiveProxy:
    """
    Proxy for yt-dlp's self.archive that delegates to our Archiver.

    Supports membership checks (`id in proxy`) and `.add(id)`.
    """

    def __init__(self, file: str | None):
        self._file: str | None = file
        "The archive file path."

    def __contains__(self, item: str) -> bool:
        if not self._file or not item:
            return False

        try:
            from app.features.ytdlp.archiver import Archiver

            status: bool = item in Archiver.get_instance().read(self._file, [item])
            return status
        except Exception:
            return False

    def add(self, item: str) -> bool:
        if not self._file or not item:
            return False

        try:
            from app.features.ytdlp.archiver import Archiver

            status: bool = Archiver.get_instance().add(self._file, [item])
            return status
        except Exception:
            return False

    def __bool__(self) -> bool:
        return bool(self._file)


class YTDLP(yt_dlp.YoutubeDL):
    _interrupted = False
    _registered = False

    def __init__(self, params=None, auto_init=True):
        apply_ytdlp_patches()

        if not YTDLP._registered:
            try:
                from app.yt_dlp_plugins.extractor import generic_browser
                from app.yt_dlp_plugins.postprocessor.nfo_maker import NFOMakerPP
                from yt_dlp.postprocessor import postprocessors

                postprocessors.value.update({"NFOMakerPP": NFOMakerPP})
                YTDLP._registered = True
            except Exception:
                get_logger().exception("Failed to register yt-dlp plugins")

        # Avoid yt-dlp preloading the archive file by stripping the param first
        orig_file = None
        patched_params: dict[str, Any] | None = None
        if params is not None:
            try:
                orig_file: str | None = params.get("download_archive")
                patched_params = dict(params)
                if "download_archive" in patched_params:
                    patched_params.pop("download_archive", None)
            except Exception:
                patched_params = params

        self._ytptube_outtmpl_info: dict[str, Any] | None = None
        self._ytptube_outtmpl_cache: dict[str, Any] = {}

        super().__init__(params=patched_params, auto_init=auto_init)

        # Restore param and replace upstream archive set with our proxy
        if orig_file is not None:
            try:
                self.params["download_archive"] = orig_file
            except Exception:
                pass

        self.archive = _ArchiveProxy(orig_file)

    def _delete_downloaded_files(self, *args, **kwargs) -> None:
        if self._interrupted:
            self.to_screen("[info] Cancelled — skipping temp cleanup.")
            return None

        return super()._delete_downloaded_files(*args, **kwargs)

    def _reset_outtmpl_cache(self) -> None:
        self._ytptube_outtmpl_info = None
        self._ytptube_outtmpl_cache = {}

    def prepare_outtmpl(self, outtmpl, info_dict, sanitize=False, *, _exec=False):
        if self._ytptube_outtmpl_info is not info_dict:
            self._ytptube_outtmpl_info = info_dict
            self._ytptube_outtmpl_cache = {}

        outtmpl, enriched = rewrite_outtmpl(outtmpl, info_dict, cache=self._ytptube_outtmpl_cache)

        return super().prepare_outtmpl(outtmpl, enriched, sanitize=sanitize, _exec=_exec)

    def process_info(self, info_dict):
        try:
            return super().process_info(info_dict)
        finally:
            self._reset_outtmpl_cache()

    def record_download_archive(self, info_dict) -> None:
        if not self.params.get("download_archive"):
            return

        if not (archive_id := self._make_archive_id(info_dict)):
            return

        assert archive_id, "Expected non-empty archive ID"

        self.write_debug(f"Adding to archive: {archive_id}")
        self.archive.add(archive_id)
        old_archive_ids = info_dict.get("_old_archive_ids", [])
        if old_archive_ids and isinstance(old_archive_ids, list) and len(old_archive_ids) > 0:
            for old_id in old_archive_ids:
                if old_id == archive_id:
                    continue

                self.write_debug(f"Adding to archive (old id): {old_id}")
                self.archive.add(old_id)


def ytdlp_options() -> list[dict[str, Any]]:
    """
    Collect yt-dlp options and return them in a structured format.

    Returns:
        list[dict[str, Any]]: A list of dictionaries containing option flags, descriptions,

    """
    from yt_dlp.options import create_parser

    from app.features.ytdlp.utils import _DATA

    parser = create_parser()

    ignored_flags: set[str] = {
        f.strip() for group in _DATA.REMOVE_KEYS for v in group.values() for f in v.split(",") if f.strip()
    }

    def collect(opts, group: str) -> list[dict[str, Any]]:
        return [
            {
                "flags": list(getattr(opt, "_short_opts", [])) + list(getattr(opt, "_long_opts", [])),
                "description": getattr(opt, "help", None),
                "group": group,
                "ignored": any(
                    f in ignored_flags for f in getattr(opt, "_short_opts", []) + getattr(opt, "_long_opts", [])
                ),
            }
            for opt in opts
            if (getattr(opt, "_long_opts", []) and getattr(opt, "help", None))
        ]

    opts = collect(parser.option_list, "root") + [
        entry for grp in parser.option_groups for entry in collect(grp.option_list, grp.title or "other")
    ]

    for opt in opts:
        if "SUPPRESSHELP" in opt.get("description", ""):
            opt["description"] = "No description available from yt-dlp."

    return opts