grabber-ai/tools/relabel_studios.py

"""Доразметка студий в data/labels.csv без пересборки страниц.

То же, что теперь делает сборщик (tools/harvest/harvest_core.py) при переиндексации,
но по готовому labels.csv и data/raw — минуты вместо полного разбора HTML:

* nnm, строки «Перевод:» (studio_source=tech) с «(не указана)»: студия, названная
  в самой части строки («Дублированный (Невафильм)», «… - Jaskier») — nnm_part_studio;
* kinozal: студии из заголовка раздачи «ПМ (LostFilm, HDrezka Studio)», которых нет
  среди строк релиза (studio_source=title) — kz_title_rows по data/raw/kinozal/<id>.json;
* все трекеры: студии из кодек-строк «Аудио 2: AC3 … (MVO, LostFilm)», «Аудио Rus: … | TVShows»,
  которых нет среди строк релиза (studio_source=codec) — codec_line_rows по data/raw/<трекер>/<id>.json.

Ручные релизы (data/manual_ids.txt) не трогаются. Перед записью — копия
data/labels.bak_relabel_<время>.csv. Повторный запуск ничего не добавляет.

    .venv\\Scripts\\python tools\\relabel_studios.py            # применить
    .venv\\Scripts\\python tools\\relabel_studios.py --dry-run  # только статистика
"""

from __future__ import annotations

import argparse
import csv
import json
import re
import shutil
import sys
import time
from collections import Counter, defaultdict
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
sys.path.insert(0, str(ROOT / "tools" / "harvest"))

import harvest_core as H  # noqa: E402

LABELS = ROOT / "data" / "labels.csv"
RAW = ROOT / "data" / "raw"
RAW_KZ = RAW / "kinozal"
MANUAL = ROOT / "data" / "manual_ids.txt"
NO_STUDIO = "(не указана)"


def relabel(rows: list[dict], raw_kz: Path = RAW_KZ, manual: set[str] | None = None) -> tuple[list[dict], Counter]:
    """rows -> (новые rows, статистика). Исходные dict меняются на месте."""
    manual = manual or set()
    stats: Counter = Counter()
    for r in rows:
        if r["source"] == "nnm" and r.get("studio_source") == "tech" and r["studio"] == NO_STUDIO \
                and r["release_id"] not in manual:
            studio = H.nnm_part_studio(r.get("role_raw") or "")
            if studio:
                r["studio"] = studio
                stats["nnm: студия из «Перевод:»"] += 1
    by_rel: dict[str, list[dict]] = defaultdict(list)
    for r in rows:
        if r["source"] == "kinozal":
            by_rel[r["release_id"]].append(r)
    added: list[dict] = []
    for rid, rel_rows in by_rel.items():
        if rid in manual:
            continue
        path = raw_kz / f"{rid.removeprefix('kz-')}.json"
        if not path.exists():
            stats["kinozal: нет raw JSON"] += 1
            continue
        rec = json.loads(path.read_text(encoding="utf-8"))
        rec["id"] = rid
        text = "\n".join(x for x in (rec.get("release_tab") or "", rec.get("tech_tab") or "") if x)
        new = H.kz_title_rows(rec, text, rel_rows)
        if new:
            stats["kinozal: релизов со студиями из заголовка"] += 1
            stats["kinozal: студий из заголовка"] += len(new)
            stats["kinozal: из них со строкой «Аудио»"] += sum(bool(x["audio_raw"]) for x in new)
            added += new
    rows = rows + added
    by_rel = defaultdict(list)
    for r in rows:
        by_rel[(r["source"], r["release_id"])].append(r)
    codec: list[dict] = []
    for (src, rid), rel_rows in by_rel.items():
        if rid in manual:
            continue
        text = raw_text(raw_kz.parent, src, rid)
        if text is None:
            continue
        new = H.codec_line_rows({"id": rid, "source": src}, text, rel_rows)
        if new:
            stats[f"{src}: релизов со студиями из кодек-строк"] += 1
            stats[f"{src}: студий из кодек-строк"] += len(new)
            codec += new
    return rows + codec, stats


def raw_text(raw_dir: Path, src: str, rid: str) -> str | None:
    """Текст страницы, как у сборщика: kinozal — «Релиз» + «Техданные», nnm/rutracker — пост."""
    num = re.sub(r"^(kz|nnm|rt)-", "", rid)
    for name in (f"{num}.json", f"{num}.legacy.json"):
        path = raw_dir / src / name
        if path.exists():
            rec = json.loads(path.read_text(encoding="utf-8"))
            if src == "kinozal":
                return "\n".join(x for x in (rec.get("release_tab") or "", rec.get("tech_tab") or "") if x)
            return rec.get("post") or ""
    return None


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--dry-run", action="store_true", help="только посчитать, файл не менять")
    args = ap.parse_args()
    with LABELS.open(encoding="utf-8-sig", newline="") as f:
        reader = csv.DictReader(f)
        fields = reader.fieldnames
        rows = list(reader)
    manual = set(MANUAL.read_text(encoding="utf-8").split()) if MANUAL.exists() else set()
    before = sum(r["studio"] != NO_STUDIO for r in rows)
    rows, stats = relabel(rows, manual=manual)
    after = sum(r["studio"] != NO_STUDIO for r in rows)
    for k, v in stats.items():
        print(f"{k}: {v}")
    print(f"строк: {len(rows)}; со студией: {before} -> {after}")
    if args.dry_run:
        return
    backup = LABELS.with_name(f"labels.bak_relabel_{time.strftime('%Y%m%d_%H%M%S')}.csv")
    shutil.copy(LABELS, backup)
    with LABELS.open("w", encoding="utf-8-sig", newline="") as f:
        w = csv.DictWriter(f, fieldnames=fields, extrasaction="ignore")
        w.writeheader()
        w.writerows(rows)
    print(f"записано: {LABELS.relative_to(ROOT)}; копия: {backup.name}")


if __name__ == "__main__":
    main()