kba-cli 0.3.0rc4__py3-none-any.whl

kba/extract.py

@@ -0,0 +1,511 @@
+"""Read-only text extraction for PDF/Office into .txt sidecars.
+
+Sources are never modified. Sidecars are UTF-8 text at ``<out>/<rel>.txt`` so a
+normal ``kba audit start --extract-dir <out>`` picks them up without changes.
+"""
+from __future__ import annotations
+
+import hashlib
+import importlib
+import json
+import logging
+import re
+import time
+from collections import Counter
+from concurrent.futures import ThreadPoolExecutor
+from datetime import datetime
+from pathlib import Path
+
+from . import util
+
+logging.getLogger("pypdf").setLevel(logging.ERROR)
+
+SCHEMA = "kba.extract/1"
+EXTRACT_EXTS = util.EXTRACTABLE_EXTS
+LEGACY_UNSUPPORTED = util.UNSUPPORTED_TEXT_EXTS
+
+_FORMULA_STR = re.compile(r'"(?:[^"]|"")*"')
+_FORMULA_CELL = re.compile(r"\$?[A-Za-z]{1,3}\$?\d+")
+_FORMULA_NUM = re.compile(r"(?<![A-Za-z_])\d+(?:\.\d+)?")
+
+
+def _load(name: str):
+    try:
+        return importlib.import_module(name)
+    except ImportError:
+        return None
+
+
+def _missing(dep: str) -> dict:
+    return {"status": "failed", "text": "", "error": f"missing_dependency:{dep}", "xlsx_fp": None}
+
+
+def _failed(exc: Exception | str) -> dict:
+    reason = exc if isinstance(exc, str) else f"{type(exc).__name__}: {exc}"
+    return {"status": "failed", "text": "", "error": reason[:300], "xlsx_fp": None}
+
+
+def _finish(text: str, xlsx_fp: dict | None = None) -> dict:
+    text = (text or "").replace("\x00", "").strip()
+    if text:
+        return {"status": "extracted", "text": text, "error": "", "xlsx_fp": xlsx_fp}
+    return {"status": "no_text", "text": "", "error": "", "xlsx_fp": xlsx_fp}
+
+
+def _normalize_formula(formula: str) -> str:
+    s = _FORMULA_STR.sub('"S"', formula)
+    s = _FORMULA_CELL.sub("R", s)
+    s = _FORMULA_NUM.sub("N", s)
+    return re.sub(r"\s+", "", s)
+
+
+def _defined_names(wb) -> list[str]:
+    try:
+        dn = wb.defined_names
+        if hasattr(dn, "keys"):
+            return sorted(str(k) for k in dn.keys())
+        return sorted(str(x) for x in dn)
+    except Exception:  # noqa: BLE001
+        return []
+
+
+def _compact_structure(sheets: list[dict], patterns: Counter, defined: list[str]) -> dict:
+    payload = {
+        "sheets": sheets,
+        "formula_cells": sum(patterns.values()),
+        "formula_patterns": dict(patterns.most_common(100)),
+        "defined_names": defined,
+    }
+    digest = hashlib.sha256(
+        json.dumps(payload, ensure_ascii=False, sort_keys=True).encode("utf-8")
+    ).hexdigest()
+    return {
+        "sha256": digest,
+        "sheets": [s["name"] for s in sheets],
+        "formula_cells": payload["formula_cells"],
+        "pattern_count": len(patterns),
+        "defined_names": defined,
+    }
+
+
+def xlsx_structure(wb) -> dict:
+    """Compact structural fingerprint: sheets, formula patterns, defined names."""
+    sheets = []
+    patterns: Counter = Counter()
+    try:
+        for ws in wb.worksheets:
+            pats: Counter = Counter()
+            for row in ws.iter_rows():
+                for cell in row:
+                    value = cell.value
+                    if isinstance(value, str) and value.startswith("="):
+                        pats[_normalize_formula(value)] += 1
+            patterns.update(pats)
+            sheets.append({
+                "name": ws.title,
+                "max_row": int(ws.max_row or 0),
+                "max_column": int(ws.max_column or 0),
+                "formula_patterns": dict(pats.most_common(50)),
+            })
+    except Exception:  # noqa: BLE001 - partial workbook is still useful
+        pass
+    return _compact_structure(sheets, patterns, _defined_names(wb))
+
+
+def _norm_edge(line: str) -> str:
+    s = re.sub(r"\d+", "#", line.strip().lower())
+    return re.sub(r"\s+", " ", s)
+
+
+def _strip_pdf_boilerplate(pages: list[str]) -> list[str]:
+    """Remove repeated page headers/footers (only edge lines of each page)."""
+    pages = [p or "" for p in pages]
+    if len(pages) < 3:
+        return pages
+    edge_counts: Counter = Counter()
+    for page in pages:
+        lines = [ln.strip() for ln in page.splitlines() if ln.strip()]
+        # Count each edge line once per page; duplicated top/bottom lines are one vote.
+        edge_lines = list(dict.fromkeys(lines[:3] + lines[-3:]))
+        for ln in edge_lines:
+            edge_counts[_norm_edge(ln)] += 1
+    threshold = max(2, int(len(pages) * 0.6 + 0.999))
+    boiler = {line for line, count in edge_counts.items()
+              if count >= threshold and len(line) >= 3}
+    if not boiler:
+        return pages
+    cleaned = []
+    for page in pages:
+        lines = page.splitlines()
+        nonempty = [i for i, ln in enumerate(lines) if ln.strip()]
+        edge = set(nonempty[:3] + nonempty[-3:])
+        keep = [ln for i, ln in enumerate(lines)
+                if not (i in edge and _norm_edge(ln) in boiler)]
+        cleaned.append("\n".join(keep).strip())
+    return cleaned
+
+
+def _extract_pdf(path: Path) -> dict:
+    pypdf = _load("pypdf")
+    if pypdf is None:
+        return _missing("pypdf")
+    try:
+        reader = pypdf.PdfReader(str(path))
+        if getattr(reader, "is_encrypted", False):
+            try:
+                if reader.decrypt("") == 0:
+                    return _failed("encrypted")
+            except Exception:  # noqa: BLE001
+                return _failed("encrypted")
+        pages = []
+        for page in reader.pages:
+            try:
+                pages.append(page.extract_text() or "")
+            except Exception:  # noqa: BLE001 - keep other pages
+                pages.append("")
+        cleaned = [p.strip() for p in _strip_pdf_boilerplate(pages) if p.strip()]
+        return _finish("\f".join(cleaned))
+    except Exception as exc:  # noqa: BLE001
+        return _failed(exc)
+
+
+def _extract_docx(path: Path) -> dict:
+    docx = _load("docx")
+    if docx is None:
+        return _missing("python-docx")
+    try:
+        doc = docx.Document(str(path))
+        parts: list[str] = []
+        for p in doc.paragraphs:
+            t = (p.text or "").strip()
+            if t:
+                parts.append(t)
+        for table in doc.tables:
+            for row in table.rows:
+                cells = [(c.text or "").strip().replace("\n", " ") for c in row.cells]
+                if any(cells):
+                    parts.append("\t".join(cells))
+        for section in doc.sections:
+            for hf in (section.header, section.footer):
+                for p in hf.paragraphs:
+                    t = (p.text or "").strip()
+                    if t:
+                        parts.append(t)
+        return _finish("\n\n".join(parts))
+    except Exception as exc:  # noqa: BLE001
+        return _failed(exc)
+
+
+def _extract_pptx(path: Path) -> dict:
+    pptx = _load("pptx")
+    if pptx is None:
+        return _missing("python-pptx")
+    try:
+        prs = pptx.Presentation(str(path))
+        slides: list[str] = []
+        for slide in prs.slides:
+            parts: list[str] = []
+            for shape in slide.shapes:
+                if getattr(shape, "has_text_frame", False):
+                    t = (shape.text_frame.text or "").strip()
+                    if t:
+                        parts.append(t)
+                if getattr(shape, "has_table", False):
+                    for row in shape.table.rows:
+                        cells = [(c.text or "").strip().replace("\n", " ") for c in row.cells]
+                        if any(cells):
+                            parts.append("\t".join(cells))
+            try:
+                notes = slide.notes_slide.notes_text_frame.text.strip()
+                if notes:
+                    parts.append(notes)
+            except Exception:  # noqa: BLE001 - notes are optional
+                pass
+            slides.append("\n\n".join(parts))
+        return _finish("\f".join(s for s in slides if s.strip()))
+    except Exception as exc:  # noqa: BLE001
+        return _failed(exc)
+
+
+def _extract_xlsx(path: Path) -> dict:
+    openpyxl = _load("openpyxl")
+    if openpyxl is None:
+        return _missing("openpyxl")
+    try:
+        wb = openpyxl.load_workbook(str(path), read_only=True, data_only=False)
+    except Exception as exc:  # noqa: BLE001
+        return _failed(exc)
+    try:
+        parts: list[str] = []
+        sheets: list[dict] = []
+        patterns: Counter = Counter()
+        for ws in wb.worksheets:
+            parts.append(f"# sheet: {ws.title}")
+            pats: Counter = Counter()
+            for row in ws.iter_rows(values_only=True):
+                cells = ["" if v is None else str(v) for v in row]
+                for value in row:
+                    if isinstance(value, str) and value.startswith("="):
+                        pats[_normalize_formula(value)] += 1
+                while cells and cells[-1] == "":
+                    cells.pop()
+                if any(cells):
+                    parts.append("\t".join(cells))
+            patterns.update(pats)
+            sheets.append({
+                "name": ws.title,
+                "max_row": int(ws.max_row or 0),
+                "max_column": int(ws.max_column or 0),
+                "formula_patterns": dict(pats.most_common(50)),
+            })
+        xlsx_fp = _compact_structure(sheets, patterns, _defined_names(wb))
+        return _finish("\n".join(parts), xlsx_fp=xlsx_fp)
+    except Exception as exc:  # noqa: BLE001
+        return _failed(exc)
+    finally:
+        try:
+            wb.close()
+        except Exception:  # noqa: BLE001
+            pass
+
+
+def extract_one(path: Path) -> dict:
+    """Extract one file. Never raises; returns status/text/error/xlsx_fp."""
+    ext = Path(path).suffix.lower()
+    if ext in LEGACY_UNSUPPORTED:
+        return {"status": "unsupported", "text": "", "error": "legacy_office_unsupported", "xlsx_fp": None}
+    if ext == ".pdf":
+        return _extract_pdf(Path(path))
+    if ext == ".docx":
+        return _extract_docx(Path(path))
+    if ext == ".pptx":
+        return _extract_pptx(Path(path))
+    if ext == ".xlsx":
+        return _extract_xlsx(Path(path))
+    return {"status": "unsupported", "text": "", "error": "unsupported_ext", "xlsx_fp": None}
+
+
+def _out_guard(root: Path, out: Path) -> None:
+    if out == root:
+        raise ValueError("输出目录不能等于扫描根目录")
+    try:
+        rel = out.relative_to(root)
+    except ValueError:
+        return
+    if not rel.parts or rel.parts[0] not in util.DEFAULT_EXCLUDE:
+        raise ValueError(
+            "输出目录不能位于扫描根目录内（除非放在 _extract 等默认排除目录），"
+            "否则下次扫描会把 sidecar 当作文档。"
+        )
+
+
+def _process_one(p: Path, root: Path, out: Path, prior: dict,
+                 overwrite: bool, max_bytes: int) -> tuple[dict, str]:
+    ext = p.suffix.lower()
+    rel = util.rel_posix(p, root)
+    target = Path(str(out / rel) + ".txt")
+    status = ""
+    error = ""
+    xlsx_fp = None
+    text = ""
+    chars = 0
+    if ext in LEGACY_UNSUPPORTED:
+        status = "unsupported"
+        error = "legacy_office_unsupported"
+    elif target.exists() and not overwrite:
+        status = "skipped_existing"
+        old = prior.get(rel) or {}
+        if isinstance(old.get("xlsx_fp"), dict):
+            xlsx_fp = old["xlsx_fp"]
+        chars = int(old.get("chars") or 0)
+    else:
+        try:
+            too_large = p.stat().st_size > max_bytes
+        except OSError as exc:
+            too_large = False
+            status = "failed"
+            error = f"{type(exc).__name__}: {exc}"[:300]
+        if not status and too_large:
+            status = "failed"
+            error = "too_large"
+        if not status:
+            try:
+                result = extract_one(p)
+            except Exception as exc:  # noqa: BLE001
+                result = {"status": "failed", "text": "", "error": str(exc), "xlsx_fp": None}
+            status = result["status"]
+            error = result.get("error") or ""
+            xlsx_fp = result.get("xlsx_fp")
+            text = result.get("text") or ""
+            chars = len(text) if status == "extracted" else 0
+            if status in ("extracted", "no_text"):
+                try:
+                    target.parent.mkdir(parents=True, exist_ok=True)
+                    target.write_text(text if status == "extracted" else "", encoding="utf-8")
+                except OSError as exc:
+                    status = "failed"
+                    error = f"write_error: {exc}"[:300]
+                    chars = 0
+            else:
+                try:
+                    if target.exists():
+                        target.unlink()
+                except OSError:
+                    pass
+    entry = {
+        "p": rel,
+        "ext": ext,
+        "status": status,
+        "chars": chars,
+        "error": error,
+        "xlsx_fp": xlsx_fp,
+    }
+    missing = error.split(":", 1)[1] if status == "failed" and error.startswith("missing_dependency:") else ""
+    return entry, missing
+
+
+def extract_root(
+    root: Path,
+    out: Path,
+    exclude=None,
+    overwrite: bool = False,
+    max_file_mb: int = 200,
+    progress=None,
+    workers: int = 1,
+    ttl_days: int | None = None,
+) -> dict:
+    root = Path(root).resolve()
+    out = Path(out).resolve()
+    if not root.is_dir():
+        raise NotADirectoryError(str(root))
+    _out_guard(root, out)
+    out.mkdir(parents=True, exist_ok=True)
+    if ttl_days is not None:
+        purge_sidecars(out, older_than_days=ttl_days)
+    prior = read_manifest(out)
+
+    max_bytes = max_file_mb * 1024 * 1024
+    files = [p for p in util.iter_files(root, exclude=exclude)
+             if p.suffix.lower() in (EXTRACT_EXTS | LEGACY_UNSUPPORTED)]
+    counts: Counter = Counter()
+    missing: Counter = Counter()
+    entries: list[dict] = []
+    workers = max(1, min(int(workers or 1), 16))
+
+    def consume(index: int, entry: dict, missing_dep: str):
+        counts[entry["status"]] += 1
+        if missing_dep:
+            missing[missing_dep] += 1
+        entries.append(entry)
+        if progress and index % 20 == 0:
+            progress(index, len(files))
+
+    if workers > 1 and len(files) > 1:
+        with ThreadPoolExecutor(max_workers=workers) as pool:
+            futures = [pool.submit(_process_one, p, root, out, prior, overwrite, max_bytes)
+                       for p in files]
+            for i, future in enumerate(futures, 1):
+                entry, missing_dep = future.result()
+                consume(i, entry, missing_dep)
+    else:
+        for i, p in enumerate(files, 1):
+            entry, missing_dep = _process_one(p, root, out, prior, overwrite, max_bytes)
+            consume(i, entry, missing_dep)
+
+    manifest = {
+        "schema": SCHEMA,
+        "created_at": datetime.now().astimezone().isoformat(timespec="seconds"),
+        "root": {"name": root.name},
+        "params": {"overwrite": bool(overwrite), "max_file_mb": max_file_mb,
+                   "workers": workers, "ttl_days": ttl_days},
+        "summary": {
+            "files": len(files),
+            "extracted": counts.get("extracted", 0),
+            "no_text": counts.get("no_text", 0),
+            "failed": counts.get("failed", 0),
+            "unsupported": counts.get("unsupported", 0),
+            "skipped_existing": counts.get("skipped_existing", 0),
+            "missing_dependencies": dict(sorted(missing.items())),
+        },
+        "files": entries,
+    }
+    util.write_json(out / "extract.json", manifest)
+    return manifest
+
+
+def read_manifest(extract_dir: Path) -> dict:
+    """Return {rel: entry} from <extract_dir>/extract.json (tolerates absence)."""
+    data = util.read_json(Path(extract_dir) / "extract.json", {}) or {}
+    entries = data.get("files") if isinstance(data, dict) else None
+    result = {}
+    if isinstance(entries, list):
+        for entry in entries:
+            if isinstance(entry, dict) and entry.get("p"):
+                result[str(entry["p"])] = entry
+    return result
+
+
+def purge_sidecars(extract_dir: Path, older_than_days: int | None = None) -> dict:
+    """Delete sidecars listed in extract.json, never touching source files.
+
+    older_than_days=None purges all; otherwise only files older than N days.
+    """
+    out = Path(extract_dir).resolve()
+    manifest_path = out / "extract.json"
+    data = util.read_json(manifest_path, {}) or {}
+    entries = data.get("files") if isinstance(data, dict) else None
+    if not isinstance(entries, list):
+        return {"files": 0, "remaining": 0, "manifest": False}
+
+    cutoff = None if older_than_days is None else time.time() - max(0, int(older_than_days)) * 86400
+    remaining: list[dict] = []
+    purged = 0
+    for entry in entries:
+        if not isinstance(entry, dict) or not entry.get("p"):
+            continue
+        target = (out / (str(entry["p"]) + ".txt")).resolve()
+        try:
+            target.relative_to(out)
+        except ValueError:
+            continue
+        try:
+            if not target.exists():
+                continue
+            if cutoff is not None and target.stat().st_mtime >= cutoff:
+                remaining.append(entry)
+                continue
+            target.unlink()
+            purged += 1
+        except OSError:
+            remaining.append(entry)
+
+    for d in sorted((d for d in out.rglob("*") if d.is_dir()),
+                    key=lambda d: len(d.parts), reverse=True):
+        try:
+            if not any(d.iterdir()):
+                d.rmdir()
+        except OSError:
+            pass
+
+    counts = Counter(str(e.get("status") or "") for e in remaining)
+    summary = data.get("summary") if isinstance(data.get("summary"), dict) else {}
+    summary.update({
+        "files": len(remaining),
+        "extracted": counts.get("extracted", 0),
+        "no_text": counts.get("no_text", 0),
+        "failed": counts.get("failed", 0),
+        "unsupported": counts.get("unsupported", 0),
+        "skipped_existing": counts.get("skipped_existing", 0),
+        "purged": int(summary.get("purged", 0) or 0) + purged,
+    })
+    data["files"] = remaining
+    data["summary"] = summary
+    data["purged_at"] = datetime.now().astimezone().isoformat(timespec="seconds")
+    if remaining:
+        util.write_json(manifest_path, data)
+    else:
+        try:
+            manifest_path.unlink()
+        except OSError:
+            pass
+    return {"files": purged, "remaining": len(remaining), "manifest": bool(remaining)}