catalog_doc/catalogizer/scanner.py

"""Обход библиотеки и наполнение базы (этап «scan»)."""
from __future__ import annotations

import json
import os
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime
from pathlib import Path

from . import config, db
from .extract import extract, long_path
from .summarize import analyze_language, doc_terms, guess_kind
from .textutil import pretty_name

MAX_TRIES = 3  # после стольких обрывов файл помечается как непригодный


def iter_files(root: Path, extensions: set[str] | None = None):
    """Все файлы-документы библиотеки (с фильтрацией мусора и SDK-папок)."""
    exts = extensions or config.DOC_EXTENSIONS
    for dirpath, dirnames, filenames in os.walk(long_path(root)):
        dirnames[:] = [d for d in dirnames
                       if d not in config.SKIP_DIR_NAMES and not d.startswith(".")]
        for name in filenames:
            ext = os.path.splitext(name)[1].lower()
            if ext not in exts:
                continue
            if name.lower() in config.NOISE_FILENAMES:
                continue
            full = os.path.join(dirpath, name)
            try:
                st = os.stat(full)
            except OSError:
                continue
            if st.st_size < config.MIN_FILE_SIZE:
                continue
            clean = full[4:] if full.startswith("\\\\?\\") else full
            yield Path(clean), st


def _base_record(path: Path, st, root: Path) -> dict:
    rel = os.path.relpath(str(path), str(root))
    return {
        "path": str(path),
        "rel_path": rel,
        "folder": os.path.dirname(rel) or ".",
        "filename": path.name,
        "ext": path.suffix.lower(),
        "size": st.st_size,
        "mtime": st.st_mtime,
        "indexed_at": datetime.now().isoformat(timespec="seconds"),
    }


def _build_record(path: Path, st, root: Path, max_pages: int) -> dict:
    res = extract(path, max_pages)
    text = res["text"]
    terms = doc_terms(text, res["title"], res["toc"])
    rec = _base_record(path, st, root)
    rec.update({
        "title": res["title"],
        "author": res["author"],
        "year": res["year"],
        "pages": res["pages"],
        "lang": analyze_language(text),
        "kind": guess_kind(res["title"], text, res["pages"]),
        "text_chars": len(text),
        "scanned": int(len(text) < config.MIN_TEXT_CHARS),
        "error": res["error"],
        "body": text,
        "toc": json.dumps(res["toc"][:120], ensure_ascii=False),
        "terms": json.dumps(dict(terms.most_common(400)), ensure_ascii=False),
    })
    return rec


def _stub_record(path: Path, st, root: Path, reason: str) -> dict:
    """Запись о файле, который разобрать не удалось: он остаётся в каталоге."""
    rec = _base_record(path, st, root)
    rec.update({
        "title": pretty_name(path.name), "author": "", "year": None,
        "pages": None, "lang": "?", "kind": "документ", "text_chars": 0,
        "scanned": 1, "error": reason, "body": "", "toc": "[]", "terms": "{}",
    })
    return rec


def scan(con, root: Path, *, max_pages: int = config.MAX_PDF_PAGES,
         jobs: int = 4, force: bool = False, limit: int | None = None,
         extensions: set[str] | None = None, quiet: bool = False,
         progress=None, stop=None) -> dict:
    """Обходит библиотеку и обновляет записи изменившихся файлов.

    Обход инкрементальный: повторный запуск читает только новые и
    изменившиеся файлы, поэтому прерванный разбор можно продолжить.
    """
    known = db.get_state(con)
    broken = db.failed_paths(con, MAX_TRIES)
    alive: set[str] = set()
    todo: list[tuple[Path, os.stat_result]] = []
    skipped = 0

    for path, st in iter_files(root, extensions):
        sp = str(path)
        alive.add(sp)
        prev = known.get(sp)
        if not force and prev and abs(prev[0] - st.st_mtime) < 1 \
                and prev[1] == st.st_size:
            continue
        if sp in broken:
            # Файл уже несколько раз обрывал разбор — заносим без текста.
            db.upsert_doc(con, _stub_record(
                path, st, root,
                "разбор аварийно прерывался %d раз(а)" % MAX_TRIES))
            skipped += 1
            continue
        todo.append((path, st))
        if limit and len(todo) >= limit:
            break

    removed = 0 if limit else db.purge_missing(con, alive)
    con.commit()
    total = len(todo)
    msg = ("Найдено файлов: %d, к разбору: %d, пропущено битых: %d, "
           "удалено из базы: %d" % (len(alive), total, skipped, removed))
    if not quiet:
        print(msg, flush=True)
    if progress:
        progress(0, total, msg)

    done = errors = 0
    started = time.time()
    chunk = max(jobs, 10)
    for offset in range(0, total, chunk):
        batch = todo[offset: offset + chunk]
        # Попытка отмечается до разбора: если процесс погибнет внутри
        # библиотеки-парсера, при следующем запуске файл будет опознан.
        for p, _ in batch:
            db.bump_attempt(con, str(p))
        con.commit()

        with ThreadPoolExecutor(max_workers=jobs) as pool:
            futures = {pool.submit(_build_record, p, s, root, max_pages): (p, s)
                       for p, s in batch}
            for fut in as_completed(futures):
                path, st = futures[fut]
                try:
                    rec = fut.result()
                except Exception as e:  # noqa: BLE001
                    rec = _stub_record(path, st, root,
                                       "%s: %s" % (type(e).__name__, e))
                db.upsert_doc(con, rec)
                db.clear_attempt(con, str(path))
                done += 1
                if rec["error"]:
                    errors += 1
        con.commit()
        speed = done / max(time.time() - started, 0.1)
        eta = (total - done) / max(speed, 0.001)
        line = ("%d/%d  (%.1f файл/с, осталось ~%d мин)"
                % (done, total, speed, eta / 60))
        if not quiet:
            print("  " + line, flush=True)
        if progress:
            progress(done, total, line)
        if stop and stop():
            break

    db.set_meta(con, "root", str(root))
    db.set_meta(con, "last_scan", datetime.now().isoformat(timespec="seconds"))
    con.commit()
    return {"total_files": len(alive), "processed": done, "errors": errors,
            "removed": removed, "skipped": skipped,
            "seconds": round(time.time() - started, 1)}