SKYNET

143 lignes · 4,5 Ko

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
extracteur.py — lit le texte d'un CV (PDF, DOCX ou texte) et l'écrit en JSON sur la sortie standard.

    python3 extracteur.py /chemin/fichier [--ocr]

Lancé par app.py dans un processus séparé : si un fichier piégé fait boucler ou saturer une bibliothèque,
seul ce processus est arrêté (limites de mémoire et de temps ci-dessous), pas le service.
Sortie : {"texte": "...", "methode": "pypdf|pdfminer|ocr|docx|texte|aucune", "pages": N}
"""
import json
import os
import re
import sys
import zipfile

PAGES_MAX = 12
OCR_PAGES_MAX = 3
TEXTE_MAX = 200_000
LETTRES_MIN = 60
DOCX_MAX_OCTETS = 8 * 1024 * 1024


def limiter_ressources():
    try:
        import resource
        resource.setrlimit(resource.RLIMIT_AS, (2 * 1024 ** 3, 2 * 1024 ** 3))
        resource.setrlimit(resource.RLIMIT_CPU, (50, 50))
    except Exception:
        pass


def nb_lettres(texte):
    return sum(1 for c in texte if c.isalpha())


def assez_de_texte(texte):
    """Du vrai texte, pas une bouillie de symboles : assez de lettres et elles dominent."""
    signes = sum(1 for c in texte if not c.isspace())
    lettres = nb_lettres(texte)
    return lettres >= LETTRES_MIN and signes > 0 and lettres / signes >= 0.6


def lire_pypdf(chemin):
    from pypdf import PdfReader
    lecteur = PdfReader(chemin, strict=False)
    if lecteur.is_encrypted:
        try:
            lecteur.decrypt("")
        except Exception:
            return "", 0
    textes = []
    pages = lecteur.pages[:PAGES_MAX]
    for page in pages:
        try:
            textes.append(page.extract_text() or "")
        except Exception:
            textes.append("")
    return "\n".join(textes), len(pages)


def lire_pdfminer(chemin):
    from pdfminer.high_level import extract_text
    return extract_text(chemin, maxpages=PAGES_MAX) or ""


def lire_ocr(chemin):
    """CV scanné (image) : reconnaissance de texte sur les premières pages. Demande tesseract + poppler."""
    from pdf2image import convert_from_path
    import pytesseract
    images = convert_from_path(chemin, dpi=200, first_page=1, last_page=OCR_PAGES_MAX)
    return "\n".join(pytesseract.image_to_string(image, lang=os.environ.get("CV_OCR_LANG", "fra+eng")) for image in images)


def lire_docx(chemin):
    with zipfile.ZipFile(chemin) as z:
        info = z.getinfo("word/document.xml")
        if info.file_size > DOCX_MAX_OCTETS:
            return ""
        xml = z.read("word/document.xml").decode("utf-8", errors="replace")
    xml = re.sub(r"</w:p>", "\n", xml)
    xml = re.sub(r"<w:(?:tab|br)\s*/>", " ", xml)
    xml = re.sub(r"<[^>]+>", "", xml)
    return (xml.replace("&amp;", "&").replace("&lt;", "<").replace("&gt;", ">")
               .replace("&quot;", '"').replace("&apos;", "'"))


def main():
    if len(sys.argv) < 2:
        sys.stderr.write("usage : extracteur.py fichier [--ocr]\n")
        return 2
    limiter_ressources()
    chemin = sys.argv[1]
    ocr = "--ocr" in sys.argv[2:]
    with open(chemin, "rb") as f:
        debut = f.read(8)

    texte, methode, pages = "", "aucune", 0
    if debut.startswith(b"%PDF"):
        for nom, fonction in (("pypdf", None), ("pdfminer", lire_pdfminer)):
            try:
                if nom == "pypdf":
                    t, pages = lire_pypdf(chemin)
                else:
                    t = fonction(chemin)
            except ImportError:
                continue
            except Exception as erreur:
                sys.stderr.write("%s : %s\n" % (nom, erreur))
                continue
            if len(t) > len(texte):
                texte, methode = t, nom
            if assez_de_texte(t):
                texte, methode = t, nom
                break
        if not assez_de_texte(texte) and ocr:
            try:
                t = lire_ocr(chemin)
                if nb_lettres(t) > nb_lettres(texte):
                    texte, methode = t, "ocr"
            except Exception as erreur:
                sys.stderr.write("ocr : %s\n" % erreur)
    elif debut.startswith(b"PK"):
        try:
            texte, methode = lire_docx(chemin), "docx"
        except Exception as erreur:
            sys.stderr.write("docx : %s\n" % erreur)
    else:
        brut = open(chemin, "rb").read(TEXTE_MAX * 4)
        try:
            texte = brut.decode("utf-8")
        except UnicodeDecodeError:
            texte = brut.decode("cp1252", errors="replace")
        methode = "texte"

    json.dump({"texte": texte[:TEXTE_MAX], "methode": methode, "pages": pages}, sys.stdout, ensure_ascii=False)
    return 0


if __name__ == "__main__":
    sys.exit(main())