douyin_video_download/douyin/download.py

"""断点续传下载 + 已下载去重。"""

from __future__ import annotations

import json
import os
import re
import threading
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, Iterable, Optional

import requests

CHUNK = 256 * 1024

# Windows 文件名非法字符，外加控制字符
_ILLEGAL = re.compile(r'[\\/:*?"<>|\x00-\x1f]')


def safe_filename(name: str, max_len: int = 80) -> str:
    name = _ILLEGAL.sub("_", name).strip(" .")
    name = re.sub(r"\s+", " ", name)
    if len(name) > max_len:
        name = name[:max_len].rstrip()
    return name or "untitled"


class Registry:
    """记录已下载的 aweme_id，避免重复下载。

    批量场景下会被多线程写，所以加锁；并且不再每条都落盘（5000 条时
    每次全量 dump 是 O(n^2) 的 IO），改成攒够 save_every 条再写，
    结束时用 flush() 收尾。文件名本身就是 aweme_id，记录丢了也能靠
    目录重建，所以这个折中是安全的。
    """

    def __init__(self, path: Path, save_every: int = 1) -> None:
        self.path = path
        self.save_every = save_every
        self._lock = threading.Lock()
        self._pending = 0
        self._data: Dict[str, dict] = {}
        if path.exists():
            try:
                self._data = json.loads(path.read_text(encoding="utf-8"))
            except (ValueError, OSError):
                self._data = {}

    def __contains__(self, aweme_id: str) -> bool:
        entry = self._data.get(aweme_id)
        if not entry:
            return False
        # 记录在但文件被手动删了，就当没下过
        return Path(entry.get("file", "")).exists()

    def get(self, aweme_id: str) -> Optional[dict]:
        return self._data.get(aweme_id)

    def add(self, aweme_id: str, file: Path, size: int, desc: str = "",
            author: str = "", uri: str = "") -> None:
        """记一条。文件名只是 aweme_id，作者/标题这些就全靠这里留存。"""
        with self._lock:
            self._data[aweme_id] = {
                "file": str(file),
                "size": size,
                "author": author,
                "desc": desc,
                "uri": uri,
                "time": int(time.time()),
            }
            self._pending += 1
            if self._pending >= self.save_every:
                self._save_locked()

    def flush(self) -> None:
        """强制落盘。批量结束时调用。"""
        with self._lock:
            if self._pending:
                self._save_locked()

    def save(self) -> None:
        with self._lock:
            self._save_locked()

    def _save_locked(self) -> None:
        self.path.parent.mkdir(parents=True, exist_ok=True)
        tmp = self.path.with_suffix(self.path.suffix + ".tmp")
        tmp.write_text(json.dumps(self._data, ensure_ascii=False, indent=2),
                       encoding="utf-8")
        tmp.replace(self.path)
        self._pending = 0


@dataclass
class DownloadResult:
    path: Path
    size: int
    resumed: bool
    skipped: bool = False


def _human(n: float) -> str:
    for unit in ("B", "KB", "MB", "GB"):
        if n < 1024:
            return f"{n:.1f}{unit}"
        n /= 1024
    return f"{n:.1f}TB"


def download(session: requests.Session, urls: Iterable[str], dest: Path,
             timeout: float = 30.0, retries: int = 4,
             progress: bool = True) -> DownloadResult:
    """下载到 dest，支持 Range 续传；urls 为多个备用直链，逐个尝试。

    抖音的视频直链会 302 到 CDN，且校验 Referer/UA，所以必须带上原 session 的头。
    """
    dest.parent.mkdir(parents=True, exist_ok=True)
    part = dest.with_suffix(dest.suffix + ".part")

    if dest.exists() and dest.stat().st_size > 0:
        return DownloadResult(dest, dest.stat().st_size, resumed=False, skipped=True)

    url_list = [u for u in urls if u]
    if not url_list:
        raise ValueError("没有可用的下载直链")

    last_err: Optional[Exception] = None
    for url in url_list:
        for attempt in range(1, retries + 1):
            done = part.stat().st_size if part.exists() else 0
            headers = {"Range": f"bytes={done}-"} if done else {}
            try:
                with session.get(url, headers=headers, stream=True,
                                 timeout=timeout, allow_redirects=True) as r:
                    if done and r.status_code == 200:
                        # 服务端不认 Range，只能从头来
                        done = 0
                        part.unlink(missing_ok=True)
                    elif done and r.status_code != 206:
                        r.raise_for_status()
                    else:
                        r.raise_for_status()

                    total = int(r.headers.get("Content-Length") or 0) + done
                    mode = "ab" if done else "wb"
                    with open(part, mode) as f:
                        for chunk in r.iter_content(CHUNK):
                            if not chunk:
                                continue
                            f.write(chunk)
                            done += len(chunk)
                            if progress:
                                pct = f"{done / total * 100:5.1f}%" if total else "  ?  "
                                print(f"\r  {pct}  {_human(done)}"
                                      + (f" / {_human(total)}" if total else ""),
                                      end="", flush=True)
                if progress:
                    print()

                if total and done < total:
                    raise IOError(f"下载不完整: {done}/{total}")

                part.replace(dest)
                return DownloadResult(dest, done, resumed=bool(headers))

            except (requests.RequestException, IOError) as exc:
                last_err = exc
                if progress:
                    print(f"\r  第 {attempt}/{retries} 次失败: {exc}")
                time.sleep(min(2 ** attempt, 8))

    raise RuntimeError(f"所有直链均下载失败，最后一个错误: {last_err}")