初始提交:抖音视频批量采集流水线
App 端免签名接口(设备注册 + detail + CDN 直链),一体化流水线 下载→抽音→传COS→记录→删本地,AI 作品四档判定,四份文档: README(使用指南)、PROGRESS(批次进度)、RE(逆向全过程)、AI(判定规则)。 Co-Authored-By: Claude Code <noreply@anthropic.com>
This commit is contained in:
174
douyin/download.py
Normal file
174
douyin/download.py
Normal file
@@ -0,0 +1,174 @@
|
||||
"""断点续传下载 + 已下载去重。"""
|
||||
|
||||
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}")
|
||||
Reference in New Issue
Block a user