Format Python code with black

This commit is contained in:
kpj
2022-02-27 15:26:41 +01:00
parent 13885968e3
commit fc695896dd
136 changed files with 6142 additions and 3811 deletions

View File

@@ -5,8 +5,7 @@ import re
from PIL import Image
from django.conf import settings
from documents.parsers import DocumentParser, ParseError, \
make_thumbnail_from_pdf
from documents.parsers import DocumentParser, ParseError, make_thumbnail_from_pdf
class NoTextFoundException(Exception):
@@ -24,7 +23,7 @@ class RasterisedDocumentParser(DocumentParser):
def extract_metadata(self, document_path, mime_type):
result = []
if mime_type == 'application/pdf':
if mime_type == "application/pdf":
import pikepdf
namespace_pattern = re.compile(r"\{(.*)\}(.*)")
@@ -37,25 +36,25 @@ class RasterisedDocumentParser(DocumentParser):
value = str(value)
try:
m = namespace_pattern.match(key)
result.append({
"namespace": m.group(1),
"prefix": meta.REVERSE_NS[m.group(1)],
"key": m.group(2),
"value": value
})
result.append(
{
"namespace": m.group(1),
"prefix": meta.REVERSE_NS[m.group(1)],
"key": m.group(2),
"value": value,
}
)
except Exception as e:
self.log(
"warning",
f"Error while reading metadata {key}: {value}. Error: "
f"{e}"
f"Error while reading metadata {key}: {value}. Error: " f"{e}",
)
return result
def get_thumbnail(self, document_path, mime_type, file_name=None):
return make_thumbnail_from_pdf(
self.archive_path or document_path,
self.tempdir,
self.logging_group)
self.archive_path or document_path, self.tempdir, self.logging_group
)
def is_image(self, mime_type):
return mime_type in [
@@ -68,17 +67,15 @@ class RasterisedDocumentParser(DocumentParser):
def has_alpha(self, image):
with Image.open(image) as im:
return im.mode in ('RGBA', 'LA')
return im.mode in ("RGBA", "LA")
def get_dpi(self, image):
try:
with Image.open(image) as im:
x, y = im.info['dpi']
x, y = im.info["dpi"]
return round(x)
except Exception as e:
self.log(
'warning',
f"Error while getting DPI from image {image}: {e}")
self.log("warning", f"Error while getting DPI from image {image}: {e}")
return None
def calculate_a4_dpi(self, image):
@@ -87,16 +84,11 @@ class RasterisedDocumentParser(DocumentParser):
width, height = im.size
# divide image width by A4 width (210mm) in inches.
dpi = int(width / (21 / 2.54))
self.log(
'debug',
f"Estimated DPI {dpi} based on image width {width}"
)
self.log("debug", f"Estimated DPI {dpi} based on image width {width}")
return dpi
except Exception as e:
self.log(
'warning',
f"Error while calculating DPI for image {image}: {e}")
self.log("warning", f"Error while calculating DPI for image {image}: {e}")
return None
def extract_text(self, sidecar_file, pdf_file):
@@ -128,60 +120,60 @@ class RasterisedDocumentParser(DocumentParser):
except Exception:
# TODO catch all for various issues with PDFminer.six.
# If PDFminer fails, fall back to OCR.
self.log("warn",
"Error while getting text from PDF document with "
"pdfminer.six", exc_info=True)
self.log(
"warn",
"Error while getting text from PDF document with " "pdfminer.six",
exc_info=True,
)
# probably not a PDF file.
return None
def construct_ocrmypdf_parameters(self,
input_file,
mime_type,
output_file,
sidecar_file,
safe_fallback=False):
def construct_ocrmypdf_parameters(
self, input_file, mime_type, output_file, sidecar_file, safe_fallback=False
):
ocrmypdf_args = {
'input_file': input_file,
'output_file': output_file,
"input_file": input_file,
"output_file": output_file,
# need to use threads, since this will be run in daemonized
# processes by django-q.
'use_threads': True,
'jobs': settings.THREADS_PER_WORKER,
'language': settings.OCR_LANGUAGE,
'output_type': settings.OCR_OUTPUT_TYPE,
'progress_bar': False
"use_threads": True,
"jobs": settings.THREADS_PER_WORKER,
"language": settings.OCR_LANGUAGE,
"output_type": settings.OCR_OUTPUT_TYPE,
"progress_bar": False,
}
if settings.OCR_MODE == 'force' or safe_fallback:
ocrmypdf_args['force_ocr'] = True
elif settings.OCR_MODE in ['skip', 'skip_noarchive']:
ocrmypdf_args['skip_text'] = True
elif settings.OCR_MODE == 'redo':
ocrmypdf_args['redo_ocr'] = True
if settings.OCR_MODE == "force" or safe_fallback:
ocrmypdf_args["force_ocr"] = True
elif settings.OCR_MODE in ["skip", "skip_noarchive"]:
ocrmypdf_args["skip_text"] = True
elif settings.OCR_MODE == "redo":
ocrmypdf_args["redo_ocr"] = True
else:
raise ParseError(
f"Invalid ocr mode: {settings.OCR_MODE}")
raise ParseError(f"Invalid ocr mode: {settings.OCR_MODE}")
if settings.OCR_CLEAN == 'clean':
ocrmypdf_args['clean'] = True
elif settings.OCR_CLEAN == 'clean-final':
if settings.OCR_MODE == 'redo':
ocrmypdf_args['clean'] = True
if settings.OCR_CLEAN == "clean":
ocrmypdf_args["clean"] = True
elif settings.OCR_CLEAN == "clean-final":
if settings.OCR_MODE == "redo":
ocrmypdf_args["clean"] = True
else:
ocrmypdf_args['clean_final'] = True
ocrmypdf_args["clean_final"] = True
if settings.OCR_DESKEW and not settings.OCR_MODE == 'redo':
ocrmypdf_args['deskew'] = True
if settings.OCR_DESKEW and not settings.OCR_MODE == "redo":
ocrmypdf_args["deskew"] = True
if settings.OCR_ROTATE_PAGES:
ocrmypdf_args['rotate_pages'] = True
ocrmypdf_args['rotate_pages_threshold'] = settings.OCR_ROTATE_PAGES_THRESHOLD # NOQA: E501
ocrmypdf_args["rotate_pages"] = True
ocrmypdf_args[
"rotate_pages_threshold"
] = settings.OCR_ROTATE_PAGES_THRESHOLD # NOQA: E501
if settings.OCR_PAGES > 0:
ocrmypdf_args['pages'] = f"1-{settings.OCR_PAGES}"
ocrmypdf_args["pages"] = f"1-{settings.OCR_PAGES}"
else:
# sidecar is incompatible with pages
ocrmypdf_args['sidecar'] = sidecar_file
ocrmypdf_args["sidecar"] = sidecar_file
if self.is_image(mime_type):
dpi = self.get_dpi(input_file)
@@ -191,29 +183,27 @@ class RasterisedDocumentParser(DocumentParser):
self.log(
"info",
f"Removing alpha layer from {input_file} "
"for compatibility with img2pdf"
"for compatibility with img2pdf",
)
with Image.open(input_file) as im:
background = Image.new('RGBA', im.size, (255, 255, 255))
background = Image.new("RGBA", im.size, (255, 255, 255))
background.alpha_composite(im)
background = background.convert('RGB')
background = background.convert("RGB")
background.save(input_file, format=im.format)
if dpi:
self.log(
"debug",
f"Detected DPI for image {input_file}: {dpi}"
)
ocrmypdf_args['image_dpi'] = dpi
self.log("debug", f"Detected DPI for image {input_file}: {dpi}")
ocrmypdf_args["image_dpi"] = dpi
elif settings.OCR_IMAGE_DPI:
ocrmypdf_args['image_dpi'] = settings.OCR_IMAGE_DPI
ocrmypdf_args["image_dpi"] = settings.OCR_IMAGE_DPI
elif a4_dpi:
ocrmypdf_args['image_dpi'] = a4_dpi
ocrmypdf_args["image_dpi"] = a4_dpi
else:
raise ParseError(
f"Cannot produce archive PDF for image {input_file}, "
f"no DPI information is present in this image and "
f"OCR_IMAGE_DPI is not set.")
f"OCR_IMAGE_DPI is not set."
)
if settings.OCR_USER_ARGS and not safe_fallback:
try:
@@ -223,13 +213,14 @@ class RasterisedDocumentParser(DocumentParser):
self.log(
"warning",
f"There is an issue with PAPERLESS_OCR_USER_ARGS, so "
f"they will not be used. Error: {e}")
f"they will not be used. Error: {e}",
)
return ocrmypdf_args
def parse(self, document_path, mime_type, file_name=None):
# This forces tesseract to use one core per page.
os.environ['OMP_THREAD_LIMIT'] = "1"
os.environ["OMP_THREAD_LIMIT"] = "1"
if mime_type == "application/pdf":
text_original = self.extract_text(None, document_path)
@@ -239,8 +230,7 @@ class RasterisedDocumentParser(DocumentParser):
original_has_text = False
if settings.OCR_MODE == "skip_noarchive" and original_has_text:
self.log("debug",
"Document has text, skipping OCRmyPDF entirely.")
self.log("debug", "Document has text, skipping OCRmyPDF entirely.")
self.text = text_original
return
@@ -251,7 +241,8 @@ class RasterisedDocumentParser(DocumentParser):
sidecar_file = os.path.join(self.tempdir, "sidecar.txt")
args = self.construct_ocrmypdf_parameters(
document_path, mime_type, archive_path, sidecar_file)
document_path, mime_type, archive_path, sidecar_file
)
try:
self.log("debug", f"Calling OCRmyPDF with args: {args}")
@@ -261,42 +252,45 @@ class RasterisedDocumentParser(DocumentParser):
self.text = self.extract_text(sidecar_file, archive_path)
if not self.text:
raise NoTextFoundException(
"No text was found in the original document")
raise NoTextFoundException("No text was found in the original document")
except EncryptedPdfError:
self.log("warning",
"This file is encrypted, OCR is impossible. Using "
"any text present in the original file.")
self.log(
"warning",
"This file is encrypted, OCR is impossible. Using "
"any text present in the original file.",
)
if original_has_text:
self.text = text_original
except (NoTextFoundException, InputFileError) as e:
self.log("warning",
f"Encountered an error while running OCR: {str(e)}. "
f"Attempting force OCR to get the text.")
self.log(
"warning",
f"Encountered an error while running OCR: {str(e)}. "
f"Attempting force OCR to get the text.",
)
archive_path_fallback = os.path.join(
self.tempdir, "archive-fallback.pdf")
sidecar_file_fallback = os.path.join(
self.tempdir, "sidecar-fallback.txt")
archive_path_fallback = os.path.join(self.tempdir, "archive-fallback.pdf")
sidecar_file_fallback = os.path.join(self.tempdir, "sidecar-fallback.txt")
# Attempt to run OCR with safe settings.
args = self.construct_ocrmypdf_parameters(
document_path, mime_type,
archive_path_fallback, sidecar_file_fallback,
safe_fallback=True
document_path,
mime_type,
archive_path_fallback,
sidecar_file_fallback,
safe_fallback=True,
)
try:
self.log("debug",
f"Fallback: Calling OCRmyPDF with args: {args}")
self.log("debug", f"Fallback: Calling OCRmyPDF with args: {args}")
ocrmypdf.ocr(**args)
# Don't return the archived file here, since this file
# is bigger and blurry due to --force-ocr.
self.text = self.extract_text(
sidecar_file_fallback, archive_path_fallback)
sidecar_file_fallback, archive_path_fallback
)
except Exception as e:
# If this fails, we have a serious issue at hand.
@@ -315,7 +309,7 @@ class RasterisedDocumentParser(DocumentParser):
self.log(
"warning",
f"No text was found in {document_path}, the content will "
f"be empty."
f"be empty.",
)
self.text = ""
@@ -325,10 +319,8 @@ def post_process_text(text):
return None
collapsed_spaces = re.sub(r"([^\S\r\n]+)", " ", text)
no_leading_whitespace = re.sub(
r"([\n\r]+)([^\S\n\r]+)", '\\1', collapsed_spaces)
no_trailing_whitespace = re.sub(
r"([^\S\n\r]+)$", '', no_leading_whitespace)
no_leading_whitespace = re.sub(r"([\n\r]+)([^\S\n\r]+)", "\\1", collapsed_spaces)
no_trailing_whitespace = re.sub(r"([^\S\n\r]+)$", "", no_leading_whitespace)
# TODO: this needs a rework
# replace \0 prevents issues with saving to postgres.