295 lines
9.4 KiB
Python
295 lines
9.4 KiB
Python
#!/usr/bin/env python3
|
||
"""单独拉取抖音数据的运行入口。
|
||
|
||
这个文件不重写采集逻辑，而是复用系统原版:
|
||
app/agents/douyin_agent.py
|
||
|
||
常用命令:
|
||
python3 scripts/douyin_login.py
|
||
python3 scripts/douyin_fetch_data.py --name "某某酒店" --city "贵阳"
|
||
python3 scripts/douyin_fetch_data.py --input hotels.txt --city "贵阳" --save-evidence
|
||
|
||
输出:
|
||
默认写到 data/douyin_probe/fetch_<时间>.json
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import asyncio
|
||
import csv
|
||
import json
|
||
import os
|
||
import sys
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
ROOT = Path(__file__).resolve().parents[1]
|
||
sys.path.insert(0, str(ROOT))
|
||
|
||
from app.agents.douyin_agent import ( # noqa: E402
|
||
_collect,
|
||
_parse_dy_comments,
|
||
_parse_dy_notes,
|
||
)
|
||
|
||
|
||
DEFAULT_OUT_DIR = ROOT / "data" / "douyin_probe"
|
||
|
||
|
||
def _ts() -> str:
|
||
return datetime.now().strftime("%Y%m%d_%H%M%S")
|
||
|
||
|
||
def _read_names(path: str | None) -> list[str]:
|
||
if not path:
|
||
return []
|
||
names: list[str] = []
|
||
for line in Path(path).read_text(encoding="utf-8").splitlines():
|
||
value = line.strip()
|
||
if value and not value.startswith("#"):
|
||
names.append(value)
|
||
return names
|
||
|
||
|
||
def _dedupe_records(records: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||
deduped: list[dict[str, Any]] = []
|
||
seen: set[tuple[str, str, str]] = set()
|
||
for record in records:
|
||
key = (
|
||
str(record.get("platform") or ""),
|
||
str(record.get("kind") or ""),
|
||
str(record.get("source_id") or record.get("url") or ""),
|
||
)
|
||
if key in seen:
|
||
continue
|
||
seen.add(key)
|
||
deduped.append(record)
|
||
return deduped
|
||
|
||
|
||
def _record_for_csv(record: dict[str, Any]) -> dict[str, Any]:
|
||
fields = [
|
||
"platform",
|
||
"kind",
|
||
"source_id",
|
||
"url",
|
||
"entity_name",
|
||
"place_natural_key",
|
||
"keyword",
|
||
"title",
|
||
"content",
|
||
"author",
|
||
"author_id",
|
||
"likes",
|
||
"comments",
|
||
"collects",
|
||
"shares",
|
||
"publish_time",
|
||
"location",
|
||
]
|
||
return {field: record.get(field, "") for field in fields}
|
||
|
||
|
||
def _write_json(path: Path, payload: dict[str, Any]) -> None:
|
||
path.parent.mkdir(parents=True, exist_ok=True)
|
||
path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
|
||
|
||
|
||
def _write_jsonl(path: Path, records: list[dict[str, Any]]) -> None:
|
||
path.parent.mkdir(parents=True, exist_ok=True)
|
||
with path.open("w", encoding="utf-8") as file:
|
||
for record in records:
|
||
file.write(json.dumps(record, ensure_ascii=False) + "\n")
|
||
|
||
|
||
def _write_csv(path: Path, records: list[dict[str, Any]]) -> None:
|
||
path.parent.mkdir(parents=True, exist_ok=True)
|
||
rows = [_record_for_csv(record) for record in records]
|
||
fieldnames = list(rows[0].keys()) if rows else list(_record_for_csv({}).keys())
|
||
with path.open("w", encoding="utf-8-sig", newline="") as file:
|
||
writer = csv.DictWriter(file, fieldnames=fieldnames)
|
||
writer.writeheader()
|
||
writer.writerows(rows)
|
||
|
||
|
||
def _build_keyword(name: str, city: str, suffix: str, keyword: str | None) -> str:
|
||
if keyword:
|
||
return keyword.strip()
|
||
parts = [city.strip(), name.strip(), suffix.strip()]
|
||
return " ".join(part for part in parts if part)
|
||
|
||
|
||
def _make_tab_picker(preferred_tab: str):
|
||
def pick(labels: list[str]) -> str | None:
|
||
if not labels:
|
||
return None
|
||
if preferred_tab and preferred_tab in labels:
|
||
return preferred_tab
|
||
if "综合" in labels:
|
||
return "综合"
|
||
return labels[0]
|
||
|
||
return pick
|
||
|
||
|
||
def fetch_one(
|
||
*,
|
||
name: str,
|
||
city: str,
|
||
suffix: str,
|
||
keyword: str | None,
|
||
tab: str,
|
||
deep: bool,
|
||
place_key: str | None,
|
||
) -> dict[str, Any]:
|
||
query = _build_keyword(name, city, suffix, keyword)
|
||
print(f"[douyin] 开始采集: name={name} keyword={query}")
|
||
result = _collect(query, _make_tab_picker(tab), deep=deep)
|
||
|
||
if result.get("logged_in") is False:
|
||
return {
|
||
"ok": False,
|
||
"need_login": True,
|
||
"name": name,
|
||
"keyword": query,
|
||
"summary": "抖音未登录。请先运行 python3 scripts/douyin_login.py 完成一次登录。",
|
||
"notes": [],
|
||
"comments": [],
|
||
"records": [],
|
||
"raw_api_count": 0,
|
||
"api_url_count": result.get("api_url_count", 0),
|
||
"tabs": result.get("tabs", []),
|
||
}
|
||
|
||
if result.get("error"):
|
||
return {
|
||
"ok": False,
|
||
"need_login": False,
|
||
"name": name,
|
||
"keyword": query,
|
||
"summary": f"采集异常: {result.get('error')}",
|
||
"notes": [],
|
||
"comments": [],
|
||
"records": [],
|
||
"raw_api_count": 0,
|
||
"api_url_count": result.get("api_url_count", 0),
|
||
"tabs": result.get("tabs", []),
|
||
}
|
||
|
||
raw_api = result.get("raw_api") or []
|
||
notes = _parse_dy_notes(raw_api, name, query)
|
||
comments = _parse_dy_comments(raw_api, name, query)
|
||
records = _dedupe_records(notes + comments)
|
||
for record in records:
|
||
if place_key:
|
||
record["place_natural_key"] = place_key
|
||
|
||
return {
|
||
"ok": True,
|
||
"need_login": False,
|
||
"name": name,
|
||
"keyword": query,
|
||
"summary": f"视频 {len(notes)} 条，评论 {len(comments)} 条，去重后 {len(records)} 条。",
|
||
"notes": notes,
|
||
"comments": comments,
|
||
"records": records,
|
||
"raw_api_count": len(raw_api),
|
||
"api_url_count": result.get("api_url_count", 0),
|
||
"tabs": result.get("tabs", []),
|
||
}
|
||
|
||
|
||
async def _save_evidence(records: list[dict[str, Any]]) -> int:
|
||
if not records:
|
||
return 0
|
||
from app.db import close_pool, init_pool, sa_save_evidence
|
||
|
||
await init_pool()
|
||
try:
|
||
return await sa_save_evidence(records)
|
||
finally:
|
||
await close_pool()
|
||
|
||
|
||
def parse_args() -> argparse.Namespace:
|
||
parser = argparse.ArgumentParser(description="按酒店/关键词单独获取抖音视频与评论数据")
|
||
parser.add_argument("--name", help="单个酒店/地点名称")
|
||
parser.add_argument("--input", help="批量名称文件，每行一个酒店/地点名称")
|
||
parser.add_argument("--city", default="贵阳", help="城市前缀，默认: 贵阳")
|
||
parser.add_argument("--suffix", default="", help="附加搜索词，例如: 酒店 入住 评价")
|
||
parser.add_argument("--keyword", help="完整搜索关键词；传入后会忽略 --city/--name/--suffix 拼接")
|
||
parser.add_argument("--tab", default="综合", help="优先选择的抖音搜索 tab，默认: 综合")
|
||
parser.add_argument("--place-key", help="写入记录的 place_natural_key，单名称时可用")
|
||
parser.add_argument("--no-deep", action="store_true", help="只采搜索页，不进入视频详情深采评论")
|
||
parser.add_argument("--save-evidence", action="store_true", help="把 records 写入系统 Evidence 表")
|
||
parser.add_argument("--out", help="JSON 输出路径")
|
||
parser.add_argument("--jsonl-out", help="JSONL records 输出路径")
|
||
parser.add_argument("--csv-out", help="CSV records 输出路径")
|
||
return parser.parse_args()
|
||
|
||
|
||
def main() -> None:
|
||
args = parse_args()
|
||
names = []
|
||
if args.name:
|
||
names.append(args.name.strip())
|
||
names.extend(_read_names(args.input))
|
||
names = [name for name in names if name]
|
||
|
||
if args.keyword and not names:
|
||
names = [args.keyword.strip()]
|
||
if not names:
|
||
raise SystemExit("请提供 --name、--input 或 --keyword")
|
||
if args.keyword and len(names) > 1:
|
||
raise SystemExit("--keyword 是完整单次搜索词，不能和多名称 --input 同时使用")
|
||
if args.place_key and len(names) > 1:
|
||
raise SystemExit("--place-key 只适合单个 --name 使用，批量请不要传")
|
||
|
||
all_results: list[dict[str, Any]] = []
|
||
all_records: list[dict[str, Any]] = []
|
||
for name in names:
|
||
result = fetch_one(
|
||
name=name,
|
||
city=args.city,
|
||
suffix=args.suffix,
|
||
keyword=args.keyword,
|
||
tab=args.tab,
|
||
deep=not args.no_deep,
|
||
place_key=args.place_key,
|
||
)
|
||
print(f"[douyin] {result['summary']}")
|
||
all_results.append(result)
|
||
all_records.extend(result.get("records") or [])
|
||
|
||
all_records = _dedupe_records(all_records)
|
||
saved = 0
|
||
if args.save_evidence:
|
||
saved = asyncio.run(_save_evidence(all_records))
|
||
print(f"[douyin] 已写入 Evidence: {saved}")
|
||
|
||
out_dir = DEFAULT_OUT_DIR
|
||
out_path = Path(args.out) if args.out else out_dir / f"fetch_{_ts()}.json"
|
||
jsonl_path = Path(args.jsonl_out) if args.jsonl_out else out_path.with_suffix(".jsonl")
|
||
csv_path = Path(args.csv_out) if args.csv_out else out_path.with_suffix(".csv")
|
||
|
||
payload = {
|
||
"ok": all(result.get("ok") for result in all_results),
|
||
"need_login": any(result.get("need_login") for result in all_results),
|
||
"saved_evidence": saved,
|
||
"total_records": len(all_records),
|
||
"results": all_results,
|
||
}
|
||
_write_json(out_path, payload)
|
||
_write_jsonl(jsonl_path, all_records)
|
||
_write_csv(csv_path, all_records)
|
||
|
||
print(f"[douyin] JSON: {out_path}")
|
||
print(f"[douyin] JSONL: {jsonl_path}")
|
||
print(f"[douyin] CSV: {csv_path}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
os.chdir(ROOT)
|
||
main()
|