autotrain/scripts/add_invoice_passenger_name.py

#!/usr/bin/env python3
"""Extract passenger names from railway invoice PDFs and add them to filenames."""

from __future__ import annotations

import argparse
import csv
import re
from datetime import datetime
from pathlib import Path

from pypdf import PdfReader


ID_NAME_RE = re.compile(r"\d{6,18}\*+\d{3}[\dXx]\s*([一-鿿·]{2,8})")
NAME_RE = re.compile(r"^[一-鿿·]{2,8}$")
KEYWORD_RE = re.compile(r"电子发票|铁路电子客票")
INVALID_FILENAME_CHARS_RE = re.compile(r"[\\/:*?\"<>|\r\n\t]")


def extract_pdf_text(path: Path) -> str:
    reader = PdfReader(str(path))
    return "\n".join(page.extract_text() or "" for page in reader.pages)


def extract_passenger_name(text: str) -> str | None:
    compact = re.sub(r"\s+", " ", text)
    match = ID_NAME_RE.search(compact)
    if match:
        return match.group(1)

    lines = [line.strip() for line in text.splitlines() if line.strip()]
    for index, line in enumerate(lines):
        if re.search(r"\d{6,18}\*+\d{3}[\dXx]", line):
            tail = re.split(r"\d{6,18}\*+\d{3}[\dXx]", line, maxsplit=1)[-1].strip()
            if NAME_RE.match(tail):
                return tail
            if index + 1 < len(lines) and NAME_RE.match(lines[index + 1]):
                return lines[index + 1]
    return None


def clean_name(name: str) -> str:
    return INVALID_FILENAME_CHARS_RE.sub("", name).strip()


def build_target(path: Path, name: str) -> Path:
    stem = path.stem
    if name in stem:
        return path

    match = KEYWORD_RE.search(stem)
    if match:
        new_stem = f"{stem[:match.start()].rstrip('_')}_{name}_{stem[match.start():].lstrip('_')}"
    else:
        new_stem = f"{stem}_{name}"

    candidate = path.with_name(f"{new_stem}{path.suffix}")
    counter = 2
    while candidate.exists() and candidate != path:
        candidate = path.with_name(f"{new_stem}_{counter}{path.suffix}")
        counter += 1
    return candidate


def main() -> int:
    parser = argparse.ArgumentParser(description="Add passenger names extracted from railway invoice PDFs to filenames.")
    parser.add_argument("directory", type=Path, help="Directory containing invoice PDFs")
    parser.add_argument("--apply", action="store_true", help="Actually rename files; default is dry-run")
    parser.add_argument("--log", type=Path, help="CSV log path; default writes into the target directory")
    args = parser.parse_args()

    directory = args.directory.expanduser().resolve()
    if not directory.is_dir():
        raise SystemExit(f"not a directory: {directory}")

    log_path = args.log or directory / f"rename_passenger_names_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv"
    rows: list[dict[str, str]] = []
    renamed = 0
    skipped = 0
    failed = 0

    for path in sorted(directory.glob("*.pdf")):
        status = ""
        name = ""
        target = path
        error = ""
        try:
            text = extract_pdf_text(path)
            extracted = extract_passenger_name(text)
            if not extracted:
                status = "failed"
                error = "name_not_found"
                failed += 1
            else:
                name = clean_name(extracted)
                target = build_target(path, name)
                if target == path:
                    status = "skipped"
                    error = "name_already_in_filename"
                    skipped += 1
                else:
                    status = "renamed" if args.apply else "dry_run"
                    if args.apply:
                        path.rename(target)
                        renamed += 1
        except Exception as exc:  # noqa: BLE001 - keep batch processing and report per file.
            status = "failed"
            error = repr(exc)
            failed += 1

        rows.append(
            {
                "status": status,
                "name": name,
                "old_path": str(path),
                "new_path": str(target),
                "error": error,
            }
        )
        print(f"{status}: {path.name} -> {target.name}" + (f" ({error})" if error else ""))

    if args.apply or rows:
        with log_path.open("w", newline="", encoding="utf-8-sig") as fh:
            writer = csv.DictWriter(fh, fieldnames=["status", "name", "old_path", "new_path", "error"])
            writer.writeheader()
            writer.writerows(rows)

    print(f"total={len(rows)} renamed={renamed} skipped={skipped} failed={failed} log={log_path}")
    if not args.apply:
        print("dry-run only; rerun with --apply to rename files")
    return 0 if failed == 0 else 1


if __name__ == "__main__":
    raise SystemExit(main())