Marty Oehme
3b4db7b6b8
Extractor is a general protocol with the PDF extraction routine now being one implementation of the protocol. Preparation for adding multiple extractors (epub,djvu, or specific progammes) in the future.
133 lines
4.5 KiB
Python
133 lines
4.5 KiB
Python
from pathlib import Path
|
|
from typing import Any, Optional
|
|
|
|
import fitz
|
|
import Levenshtein
|
|
import magic
|
|
import papis.config
|
|
import papis.logging
|
|
|
|
from papis_extract.annotation import Annotation
|
|
|
|
logger = papis.logging.get_logger(__name__)
|
|
|
|
|
|
class PdfExtractor:
|
|
def can_process(self, filename: Path) -> bool:
|
|
if not filename.is_file():
|
|
logger.error(f"File {str(filename)} not readable.")
|
|
return False
|
|
if not self._is_pdf(filename):
|
|
return False
|
|
return True
|
|
|
|
def run(self, filename: Path) -> list[Annotation]:
|
|
"""Extract annotations from a file.
|
|
|
|
Returns all readable annotations contained in the file
|
|
passed in. Only returns Highlight or Text annotations.
|
|
"""
|
|
annotations = []
|
|
with fitz.Document(filename) as doc:
|
|
for page in doc:
|
|
for annot in page.annots():
|
|
quote, note = self._retrieve_annotation_content(page, annot)
|
|
if not quote and not note:
|
|
continue
|
|
col = (
|
|
annot.colors.get("fill")
|
|
or annot.colors.get("stroke")
|
|
or (0.0, 0.0, 0.0)
|
|
)
|
|
a = Annotation(
|
|
file=str(filename),
|
|
text=quote or "",
|
|
content=note or "",
|
|
colors=col,
|
|
type=annot.type[1],
|
|
page=(page.number or 0) + 1,
|
|
)
|
|
a.tag = self._tag_from_colorname(a.colorname or "")
|
|
annotations.append(a)
|
|
logger.debug(
|
|
f"Found {len(annotations)} "
|
|
f"{'annotation' if len(annotations) == 1 else 'annotations'} for {filename}."
|
|
)
|
|
return annotations
|
|
|
|
|
|
def _is_pdf(self, fname: Path) -> bool:
|
|
"""Check if file is a pdf, using mime type."""
|
|
return magic.from_file(fname, mime=True) == "application/pdf"
|
|
|
|
|
|
|
|
def _tag_from_colorname(self, colorname: str) -> str:
|
|
color_mapping: dict[str, str] = self._getdict("tags", "plugins.extract")
|
|
if not color_mapping:
|
|
return ""
|
|
|
|
return color_mapping.get(colorname, "")
|
|
|
|
|
|
def _retrieve_annotation_content(self,
|
|
page: fitz.Page, annotation: fitz.Annot
|
|
) -> tuple[str | None, str | None]:
|
|
"""Gets the text content of an annotation.
|
|
|
|
Returns the actual content of an annotation. Sometimes
|
|
that is only the written words, sometimes that is only
|
|
annotation notes, sometimes it is both. Runs a similarity
|
|
comparison between strings to find out whether they
|
|
should both be included or are the same, using
|
|
Levenshtein distance.
|
|
"""
|
|
content = annotation.info["content"].replace("\n", " ")
|
|
written = page.get_textbox(annotation.rect).replace("\n", " ")
|
|
|
|
# highlight with selection in note
|
|
minimum_similarity = (
|
|
papis.config.getfloat("minimum_similarity_content", "plugins.extract") or 1.0
|
|
)
|
|
if Levenshtein.ratio(content, written) > minimum_similarity:
|
|
return (content, None)
|
|
# both a highlight and a note
|
|
elif content and written:
|
|
return (written, content)
|
|
# an independent note, not a highlight
|
|
elif content:
|
|
return (None, content)
|
|
# highlight with selection not in note
|
|
elif written:
|
|
return (written, None)
|
|
# just a highlight without any text
|
|
return (None, None)
|
|
|
|
|
|
# mimics the functions in papis.config.{getlist,getint,getfloat} etc.
|
|
def _getdict(self, key: str, section: Optional[str] = None) -> dict[str, str]:
|
|
"""Dict getter
|
|
|
|
:returns: A python dict
|
|
:raises SyntaxError: Whenever the parsed syntax is either not a valid
|
|
python object or a valid python dict.
|
|
"""
|
|
rawvalue: Any = papis.config.general_get(key, section=section)
|
|
if isinstance(rawvalue, dict):
|
|
return rawvalue
|
|
try:
|
|
rawvalue = eval(rawvalue)
|
|
except Exception:
|
|
raise SyntaxError(
|
|
"The key '{}' must be a valid Python object: {}".format(key, rawvalue)
|
|
)
|
|
else:
|
|
if not isinstance(rawvalue, dict):
|
|
raise SyntaxError(
|
|
"The key '{}' must be a valid Python dict. Got: {} (type {!r})".format(
|
|
key, rawvalue, type(rawvalue).__name__
|
|
)
|
|
)
|
|
|
|
return rawvalue
|