<commit_before>import os
import logging

from pdfminer.pdfinterp import PDFResourceManager, PDFPageInterpreter
from pdfminer.converter import PDFPageAggregator
from pdfminer.layout import LAParams, LTTextBox, LTTextLine, LTFigure, LTImage

from pdfminer.pdfpage import PDFPage
from pdfminer.pdfparser import PDFParser
from pdfminer.pdfdocument import PDFDocument

from extractors.util import safe_text, text_fragments
from extractors.tesseract import extract_image_data, _extract_image_page


log = logging.getLogger(__name__)

OCR_MIN_WIDTH = 200
OCR_MIN_HEIGHT = 50


def _find_objects(objects, cls):
    for lt_obj in objects:
        if isinstance(lt_obj, cls):
            yield lt_obj
        elif isinstance(lt_obj, LTFigure):
            for obj in _find_objects(lt_obj._objs, cls):
                yield obj


def _convert_page(layout, languages):
    text_content = []
    for text_obj in _find_objects(layout._objs, (LTTextBox, LTTextLine)):
        text_content.append(text_obj.get_text())

    text = text_fragments(text_content)
    if len(text) > 3:
        # TODO: invent a smarter way to decide whether to do OCR.
        return text

    for img_obj in _find_objects(layout._objs, LTImage):
        try:
            if img_obj.width < OCR_MIN_WIDTH or \
                    img_obj.height < OCR_MIN_HEIGHT:
                continue
            data = img_obj.stream.get_data()
            img_text = extract_image_data(data, languages=languages)
            text_content.append(img_text)
        except Exception as ex:
            log.debug(ex)

    return text_fragments(text_content)


def extract_pdf(path, languages=None):
    """
    Extract content from a PDF file.

    This will attempt to use pdfminer to extract textual content from
    each page. If none is found, it'll send the images through OCR.
    """
    with open(path, 'rb') as fh:
        rsrcmgr = PDFResourceManager()
        laparams = LAParams()
        device = PDFPageAggregator(rsrcmgr, laparams=laparams)
        interpreter = PDFPageInterpreter(rsrcmgr, device)
        parser = PDFParser(fh)
        doc = PDFDocument(parser, '')
        result = {'pages': []}
        if len(doc.info):
            for k, v in doc.info[-1].items():
                k = k.lower().strip()
                if k != 'pages':
                    result[k] = safe_text(v)

        if not doc.is_extractable:
            log.warning("PDF not extractable: %s", path)
            return result

        for i, page in enumerate(PDFPage.create_pages(doc)):
            text = None
            try:
                interpreter.process_page(page)
                layout = device.get_result()
                text = _convert_page(layout, languages)
            except AttributeError as ae:
                log.debug("Failed to parse PDF page: %r", ae)

            if text is None or not len(text.strip()):
                text = _extract_image_page(path, i + 1)
            result['pages'].append(text)
        device.close()
        return result
<commit_msg>return None if no root object was recognized<commit_after>import os
import logging

from pdfminer.pdfinterp import PDFResourceManager, PDFPageInterpreter
from pdfminer.converter import PDFPageAggregator
from pdfminer.layout import LAParams, LTTextBox, LTTextLine, LTFigure, LTImage

from pdfminer.pdfpage import PDFPage
from pdfminer.pdfparser import PDFParser, PDFSyntaxError
from pdfminer.pdfdocument import PDFDocument

from extractors.util import safe_text, text_fragments
from extractors.tesseract import extract_image_data, _extract_image_page


log = logging.getLogger(__name__)

OCR_MIN_WIDTH = 200
OCR_MIN_HEIGHT = 50


def _find_objects(objects, cls):
    for lt_obj in objects:
        if isinstance(lt_obj, cls):
            yield lt_obj
        elif isinstance(lt_obj, LTFigure):
            for obj in _find_objects(lt_obj._objs, cls):
                yield obj


def _convert_page(layout, languages):
    text_content = []
    for text_obj in _find_objects(layout._objs, (LTTextBox, LTTextLine)):
        text_content.append(text_obj.get_text())

    text = text_fragments(text_content)
    if len(text) > 3:
        # TODO: invent a smarter way to decide whether to do OCR.
        return text
    return None
    # for img_obj in _find_objects(layout._objs, LTImage):
    #     try:
    #         if img_obj.width < OCR_MIN_WIDTH or \
    #                 img_obj.height < OCR_MIN_HEIGHT:
    #             continue
    #         data = img_obj.stream.get_data()
    #         img_text = extract_image_data(data, languages=languages)
    #         text_content.append(img_text)
    #     except Exception as ex:
    #         log.debug(ex)
    # return text_fragments(text_content)


def extract_pdf(path, languages=None):
    """
    Extract content from a PDF file.

    This will attempt to use pdfminer to extract textual content from
    each page. If none is found, it'll send the images through OCR.
    """
    with open(path, 'rb') as fh:
        rsrcmgr = PDFResourceManager()
        laparams = LAParams()
        device = PDFPageAggregator(rsrcmgr, laparams=laparams)
        interpreter = PDFPageInterpreter(rsrcmgr, device)
        parser = PDFParser(fh)
        try:
            doc = PDFDocument(parser, '')
        except PDFSyntaxError as pse:
            if 'No /Root object!' in pse.message:
                return None
            raise

        result = {'pages': []}
        if len(doc.info):
            for k, v in doc.info[-1].items():
                k = k.lower().strip()
                if k != 'pages':
                    result[k] = safe_text(v)

        if not doc.is_extractable:
            log.warning("PDF not extractable: %s", path)
            return None

        for i, page in enumerate(PDFPage.create_pages(doc)):
            text = None
            try:
                interpreter.process_page(page)
                layout = device.get_result()
                text = _convert_page(layout, languages)
            except AttributeError as ae:
                log.debug("Failed to parse PDF page: %r", ae)

            if text is None or not len(text.strip()):
                text = _extract_image_page(path, i + 1)
            result['pages'].append(text)
        device.close()
        return result

{'fraction_non_alphanumeric': 0.08196981731532962, 'fraction_numerical': 0.0025416997617156473, 'mean_word_length': 2.30498687664042, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2593670', 'n_tokens_mistral': 1896, 'n_tokens_neox': 1857, 'n_words': 510}