diff --git a/test/example.pdf b/test/example.pdf index bc2b9cd2..bbcda4bb 100644 Binary files a/test/example.pdf and b/test/example.pdf differ diff --git a/wrolpi/files/indexers.py b/wrolpi/files/indexers.py index 5b13695a..865c2b91 100644 --- a/wrolpi/files/indexers.py +++ b/wrolpi/files/indexers.py @@ -1,3 +1,4 @@ +import dataclasses import pathlib import re from abc import ABC @@ -122,6 +123,12 @@ class TextIndexer(Indexer, ABC): return words +@dataclasses.dataclass +class PDFMetadata: + title: str = None + author: str = None + + @register_indexer('application/pdf') class PDFIndexer(Indexer, ABC): """Uses PyPDF2 to extract text from a PDF.""" @@ -129,24 +136,47 @@ class PDFIndexer(Indexer, ABC): @classmethod def create_index(cls, file): path: pathlib.Path = file.path.path if hasattr(file.path, 'path') else file.path - a = cls.get_title(file) + + file_title = super().get_title(file) + + reader = None + if PdfReader is not None: + reader = PdfReader(path) + + data = cls.get_pdf_metadata(reader, file) if path.stat().st_size > FILE_MAX_PDF_SIZE: logger.warning(f'PDF too large to fully index: {path}') - return a, None, None, None + return data.title, data.author, file_title, None words = '' try: # PDFs are complex, don't fail to create title index if text extraction fails. - words = '\n'.join(truncate_generator_bytes(cls.get_words(path), FILE_MAX_TEXT_SIZE)) + words = '\n'.join(truncate_generator_bytes(cls.get_words(reader, path), FILE_MAX_TEXT_SIZE)) except Exception as e: logger.error(f'Failed to index {path}', exc_info=e) if PYTEST: raise - return a, None, None, words + return data.title, data.author, file_title, words @classmethod - def get_words(cls, path: pathlib.Path) -> Generator[str, None, None]: + def get_pdf_metadata(cls, reader: PdfReader, file) -> PDFMetadata: + """Extract title/author from the PDF metadata.""" + path = file.path.path if hasattr(file.path, 'path') else file.path + + data = PDFMetadata() + + if PdfReader is None: + logger.error(f'Cannot get title of {path} PyPDF2 is not installed.') + data.title = super().get_title(file) + return data + + data.title = reader.metadata.title + data.author = reader.metadata.author + return data + + @classmethod + def get_words(cls, reader: PdfReader, path: pathlib.Path) -> Generator[str, None, None]: """ Reads all text facing up in a PDF. @@ -157,7 +187,6 @@ class PDFIndexer(Indexer, ABC): yield '' return - reader = PdfReader(path) for page in reader.pages: # Get all text facing up. text = page.extract_text(0) diff --git a/wrolpi/files/test/test_lib.py b/wrolpi/files/test/test_lib.py index 77be95aa..ae1ae2c2 100644 --- a/wrolpi/files/test/test_lib.py +++ b/wrolpi/files/test/test_lib.py @@ -488,10 +488,12 @@ def test_pdf_indexer(example_pdf): a_text, b_text, c_text, d_text = indexers.PDFIndexer.create_index(file) - # The file name. - assert a_text == ['example', 'pdf'] - assert b_text is None - assert c_text is None + # The title extracted from the PDF. + assert a_text == 'WROLPi Test PDF' + # The author. + assert b_text == 'roland' + # The parsed file name. + assert c_text == ['example', 'pdf'] # Both pages of the example PDF are extracted. assert d_text == 'Page one\nPage two' @@ -505,7 +507,7 @@ def test_pdf_indexer_max_size(example_pdf): a_text, b_text, c_text, d_text = indexers.PDFIndexer.create_index(file) # The file name is still indexed. - assert a_text == ['example', 'pdf'] + assert a_text == 'WROLPi Test PDF' assert d_text is None