foxygit / ytmdl Log in
commits tags

/src/ytmdl/downloader.py · 3.32 KB

raw
from __future__ import annotations

import re
from dataclasses import dataclass
from pathlib import Path
from typing import Callable

import yt_dlp

ProgressHook = Callable[[dict], None]

_ILLEGAL_DIRNAME_CHARS = re.compile(r'[<>:"/\\|?*\x00-\x1f]')


class ResolveError(Exception):
    pass


class DownloadError(Exception):
    pass


@dataclass
class ResolveResult:
    entries: list[dict]
    playlist_title: str | None = None


def sanitize_dirname(name: str, max_length: int = 150) -> str:
    """Turn an arbitrary playlist/album title into a safe directory name."""
    cleaned = _ILLEGAL_DIRNAME_CHARS.sub("_", name).strip(" .")
    cleaned = re.sub(r"\s+", " ", cleaned)
    return cleaned[:max_length] or "playlist"


def resolve(url: str) -> ResolveResult:
    """Expand a video/playlist/album URL into a flat list of {url, title} entries."""
    ydl_opts = {
        "extract_flat": "in_playlist",
        "quiet": True,
        "no_warnings": True,
        "skip_download": True,
    }
    try:
        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
            info = ydl.extract_info(url, download=False)
    except yt_dlp.utils.DownloadError as exc:
        raise ResolveError(str(exc)) from exc

    if info is None:
        raise ResolveError(f"No data returned for {url}")

    is_playlist = info.get("entries") is not None
    raw_entries = info.get("entries") if is_playlist else [info]
    playlist_title = info.get("title") if is_playlist else None

    entries: list[dict] = []
    for entry in raw_entries:
        if entry is None:
            continue
        entry_url = entry.get("url") or entry.get("webpage_url")
        if entry_url and not entry_url.startswith("http"):
            video_id = entry.get("id") or entry_url
            entry_url = f"https://www.youtube.com/watch?v={video_id}"
        if not entry_url:
            continue
        entries.append({"url": entry_url, "title": entry.get("title") or entry_url})

    if not entries:
        raise ResolveError(f"Nothing downloadable found at {url}")

    return ResolveResult(entries=entries, playlist_title=playlist_title)


def download(
    url: str,
    output_dir: Path,
    audio_format: str = "mp3",
    progress_hook: ProgressHook | None = None,
) -> Path:
    """Download a single URL as an audio file with embedded metadata/thumbnail."""
    output_dir.mkdir(parents=True, exist_ok=True)

    hooks = [progress_hook] if progress_hook else []

    ydl_opts = {
        "format": "bestaudio/best",
        "outtmpl": str(output_dir / "%(title)s.%(ext)s"),
        "postprocessors": [
            {
                "key": "FFmpegExtractAudio",
                "preferredcodec": audio_format,
                "preferredquality": "0",
            },
            {"key": "FFmpegThumbnailsConvertor", "format": "jpg"},
            {"key": "EmbedThumbnail"},
            {"key": "FFmpegMetadata", "add_metadata": True},
        ],
        "writethumbnail": True,
        "progress_hooks": hooks,
        "quiet": True,
        "no_warnings": True,
        "noprogress": True,
        "restrictfilenames": False,
    }

    try:
        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
            info = ydl.extract_info(url, download=True)
    except yt_dlp.utils.DownloadError as exc:
        raise DownloadError(str(exc)) from exc

    title = info.get("title", url)
    return output_dir / f"{title}.{audio_format}"