259 lines
8.6 KiB
Python
259 lines
8.6 KiB
Python
"""yt-dlp wrapper: download a YouTube video together with available subtitles."""
|
||
from __future__ import annotations
|
||
|
||
import logging
|
||
import re
|
||
from pathlib import Path
|
||
from typing import Optional
|
||
|
||
import yt_dlp
|
||
|
||
|
||
logger = logging.getLogger("youtube")
|
||
|
||
|
||
_VIDEO_EXTS = {".mp4", ".mkv", ".webm", ".m4a", ".mp3", ".mov", ".opus", ".ogg", ".aac"}
|
||
|
||
|
||
def expand_url(url: str) -> dict:
|
||
"""Определяет, видео это или плейлист, и возвращает список элементов.
|
||
|
||
Ссылка вида watch?v=…&list=… (видео внутри плейлиста) трактуется как весь
|
||
плейлист — именно её люди копируют из адресной строки.
|
||
|
||
Returns: {"kind": "video"|"playlist", "title": str,
|
||
"entries": [{"url": str, "title": str}, ...]}
|
||
"""
|
||
m = re.search(r"[?&]list=([A-Za-z0-9_\-]+)", url)
|
||
if m and "playlist?" not in url:
|
||
# RD… и LL/WL — авто-миксы и служебные списки, их не разворачиваем
|
||
list_id = m.group(1)
|
||
if not list_id.startswith(("RD", "LL", "WL")):
|
||
url = f"https://www.youtube.com/playlist?list={list_id}"
|
||
|
||
opts = {
|
||
"extract_flat": "in_playlist",
|
||
"quiet": True,
|
||
"no_warnings": True,
|
||
"skip_download": True,
|
||
}
|
||
with yt_dlp.YoutubeDL(opts) as ydl:
|
||
info = ydl.extract_info(url, download=False)
|
||
|
||
if info.get("_type") == "playlist":
|
||
entries = []
|
||
for e in info.get("entries") or []:
|
||
if not e:
|
||
continue
|
||
u = e.get("url") or ""
|
||
if u and not u.startswith("http"):
|
||
u = f"https://www.youtube.com/watch?v={u}"
|
||
if not u and e.get("id"):
|
||
u = f"https://www.youtube.com/watch?v={e['id']}"
|
||
if u:
|
||
entries.append({"url": u, "title": e.get("title") or u})
|
||
return {
|
||
"kind": "playlist",
|
||
"title": info.get("title") or url,
|
||
"entries": entries,
|
||
}
|
||
|
||
title = info.get("title") or url
|
||
return {"kind": "video", "title": title, "entries": [{"url": url, "title": title}]}
|
||
|
||
|
||
def download_video_with_subs(
|
||
url: str,
|
||
video_id: str,
|
||
video_dir: Path,
|
||
langs: Optional[list[str]] = None,
|
||
audio_only: bool = False,
|
||
on_progress=None,
|
||
) -> dict:
|
||
"""Download a single video, also requesting subtitles for the given languages.
|
||
|
||
audio_only: качает только аудиодорожку (для длинных лекций, когда видео
|
||
не сохраняется — в разы быстрее и меньше по объёму).
|
||
on_progress: callback(percent: int), вызывается с шагом ~5%.
|
||
|
||
Returns:
|
||
{
|
||
"title": str,
|
||
"video_path": Path,
|
||
"subtitle_path": Path | None,
|
||
"subtitle_lang": str | None,
|
||
"subtitle_kind": "manual" | "auto" | None,
|
||
}
|
||
"""
|
||
if langs is None:
|
||
langs = ["ru", "en"]
|
||
video_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
outtmpl = str(video_dir / f"{video_id}.%(ext)s")
|
||
|
||
base_opts = {
|
||
"outtmpl": outtmpl,
|
||
"quiet": True,
|
||
"no_warnings": True,
|
||
"noplaylist": True, # ссылка с list= в задаче на одно видео не тянет весь плейлист
|
||
}
|
||
if audio_only:
|
||
base_opts["format"] = "ba/b"
|
||
else:
|
||
base_opts["format"] = "bv*+ba/b"
|
||
base_opts["merge_output_format"] = "mp4"
|
||
|
||
if on_progress is not None:
|
||
state = {"last": -5}
|
||
|
||
def _hook(dd: dict):
|
||
if dd.get("status") != "downloading":
|
||
return
|
||
total = dd.get("total_bytes") or dd.get("total_bytes_estimate")
|
||
if not total:
|
||
return
|
||
pct = int(dd.get("downloaded_bytes", 0) * 100 / total)
|
||
if pct >= state["last"] + 5:
|
||
state["last"] = pct
|
||
try:
|
||
on_progress(pct)
|
||
except Exception:
|
||
pass
|
||
|
||
base_opts["progress_hooks"] = [_hook]
|
||
sub_opts = {
|
||
**base_opts,
|
||
"writesubtitles": True,
|
||
"writeautomaticsub": True,
|
||
"subtitleslangs": langs,
|
||
"subtitlesformat": "vtt",
|
||
}
|
||
|
||
# YouTube often 429s the subtitle endpoints; a failed subtitle download must not
|
||
# kill the pipeline — retry without subs and let the Vosk fallback transcribe.
|
||
try:
|
||
with yt_dlp.YoutubeDL(sub_opts) as ydl:
|
||
info = ydl.extract_info(url, download=True)
|
||
except yt_dlp.utils.DownloadError as e:
|
||
logger.warning("Download with subtitles failed (%s); retrying without subs", e)
|
||
with yt_dlp.YoutubeDL(base_opts) as ydl:
|
||
info = ydl.extract_info(url, download=True)
|
||
|
||
title = info.get("title") or video_id
|
||
|
||
# Locate the merged video file
|
||
video_path = video_dir / f"{video_id}.mp4"
|
||
if not video_path.exists():
|
||
for p in video_dir.glob(f"{video_id}.*"):
|
||
if p.suffix.lower() in _VIDEO_EXTS:
|
||
video_path = p
|
||
break
|
||
|
||
sub_path, sub_lang, sub_kind = _find_subtitle(video_dir, video_id, langs, info)
|
||
|
||
return {
|
||
"title": title,
|
||
"video_path": video_path,
|
||
"subtitle_path": sub_path,
|
||
"subtitle_lang": sub_lang,
|
||
"subtitle_kind": sub_kind,
|
||
}
|
||
|
||
|
||
def _find_subtitle(
|
||
video_dir: Path,
|
||
video_id: str,
|
||
langs: list[str],
|
||
info: dict,
|
||
) -> tuple[Optional[Path], Optional[str], Optional[str]]:
|
||
"""Pick the best subtitle file among downloaded ones, preferring manual subs."""
|
||
requested = info.get("requested_subtitles") or {}
|
||
manual_keys = set(info.get("subtitles", {}).keys())
|
||
|
||
# Build a list of (lang, kind, path) candidates from yt-dlp's reported requested_subtitles
|
||
candidates: list[tuple[str, str, Path]] = []
|
||
for lang_code, sub_info in requested.items():
|
||
filepath = sub_info.get("filepath")
|
||
if not filepath:
|
||
continue
|
||
p = Path(filepath)
|
||
if not p.exists():
|
||
continue
|
||
kind = "manual" if lang_code in manual_keys else "auto"
|
||
candidates.append((lang_code, kind, p))
|
||
|
||
# Fallback: glob the directory
|
||
if not candidates:
|
||
for p in video_dir.glob(f"{video_id}*.vtt"):
|
||
m = re.match(rf"^{re.escape(video_id)}\.([A-Za-z0-9_\-]+)\.vtt$", p.name)
|
||
lang_code = m.group(1) if m else "unknown"
|
||
candidates.append((lang_code, "auto", p))
|
||
|
||
if not candidates:
|
||
return None, None, None
|
||
|
||
def score(c: tuple[str, str, Path]) -> tuple[int, int]:
|
||
lang, kind, _ = c
|
||
try:
|
||
lang_rank = next(
|
||
i for i, l in enumerate(langs) if lang == l or lang.startswith(l + "-")
|
||
)
|
||
except StopIteration:
|
||
lang_rank = len(langs)
|
||
kind_rank = 0 if kind == "manual" else 1
|
||
return (lang_rank, kind_rank)
|
||
|
||
candidates.sort(key=score)
|
||
lang, kind, path = candidates[0]
|
||
return path, lang, kind
|
||
|
||
|
||
def vtt_to_text(vtt_content: str) -> str:
|
||
"""Convert a VTT subtitle file to a deduplicated plain-text string.
|
||
|
||
YouTube auto-captions ship as rolling cues — each new cue extends the previous.
|
||
Strategy: per cue block, keep only the last text line, then drop consecutive
|
||
duplicates and identical sentence fragments.
|
||
"""
|
||
blocks = re.split(r"\r?\n\s*\r?\n", vtt_content)
|
||
out: list[str] = []
|
||
last = ""
|
||
|
||
for block in blocks:
|
||
block = block.strip()
|
||
if not block:
|
||
continue
|
||
text_lines: list[str] = []
|
||
for raw in block.splitlines():
|
||
line = raw.strip()
|
||
if not line:
|
||
continue
|
||
if line.startswith("WEBVTT"):
|
||
continue
|
||
if line.startswith(("NOTE", "STYLE", "Kind:", "Language:", "Region:")):
|
||
continue
|
||
if "-->" in line:
|
||
continue
|
||
if re.match(r"^\d+$", line):
|
||
continue
|
||
line = re.sub(r"<[^>]+>", "", line)
|
||
line = re.sub(r"\s+", " ", line).strip()
|
||
if line:
|
||
text_lines.append(line)
|
||
if not text_lines:
|
||
continue
|
||
candidate = text_lines[-1]
|
||
if candidate == last:
|
||
continue
|
||
# Avoid the rolling-caption case where a cue strictly extends the previous one
|
||
if last and candidate.startswith(last):
|
||
out[-1] = candidate
|
||
last = candidate
|
||
continue
|
||
if last and last.startswith(candidate):
|
||
continue
|
||
out.append(candidate)
|
||
last = candidate
|
||
|
||
return " ".join(out)
|