menekis/OCRmyPDF
Public Access
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Refactor setup_pipeline to context manager

Browse Source
This commit is contained in:
James R. Barlow
2023-10-24 00:52:31 -07:00
parent ebfe008432
commit 8985c0dfe9
4 changed files with 124 additions and 118 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/ocrmypdf/_pipelines/common.py
+17 -5
View File
@@ -5,7 +5,9 @@ import json
import logging
import logging.handlers
import os
from collections.abc import Sequence
import sys
from collections.abc import Generator, Sequence
from contextlib import contextmanager
from dataclasses import dataclass
from pathlib import Path
from tempfile import mkdtemp
@@ -14,7 +16,7 @@ from typing import NamedTuple
import PIL
from ocrmypdf._concurrent import Executor, setup_executor
from ocrmypdf._jobcontext import PageContext, PdfContext
from ocrmypdf._jobcontext import PageContext, PdfContext, cleanup_working_files
from ocrmypdf._logging import PageNumberFilter
from ocrmypdf._pipeline import (
convert_to_pdfa,
@@ -153,13 +155,14 @@ def worker_init(max_pixels: int) -> None:
pikepdf_enable_mmap()
@contextmanager
def setup_pipeline(
*,
options: argparse.Namespace,
plugin_manager: OcrmypdfPluginManager | None,
api: bool = False,
work_folder: Path | None,
) -> tuple[Path, logging.FileHandler | None, Executor, OcrmypdfPluginManager]:
) -> Generator[tuple[Path, Executor, OcrmypdfPluginManager], None, None]:
# Any changes to options will not take effect for options that are already
# bound to function parameters in the pipeline. (For example
# options.input_file, options.pdf_renderer are already bound.)
@@ -179,12 +182,21 @@ def setup_pipeline(
# See https://github.com/pytest-dev/pytest/issues/5502 for why we skip this
# when pytest is running
debug_log_handler = configure_debug_logging(
Path(work_folder) / "debug.log"
work_folder / "debug.log"
) # pragma: no cover
pikepdf_enable_mmap()
executor = setup_executor(plugin_manager)
return work_folder, debug_log_handler, executor, plugin_manager
try:
yield work_folder, executor, plugin_manager
finally:
if debug_log_handler:
try:
debug_log_handler.close()
log.removeHandler(debug_log_handler)
except OSError as e:
print(e, file=sys.stderr)
cleanup_working_files(work_folder, options)
def preprocess(
src/ocrmypdf/_pipelines/hocr_to_ocr_pdf.py
+21 -25
View File
@@ -111,30 +111,26 @@ def run_hocr_to_ocr_pdf_pipeline(
*,
plugin_manager: OcrmypdfPluginManager | None,
) -> None:
work_folder, debug_log_handler, executor, plugin_manager = setup_pipeline(
options=options,
plugin_manager=plugin_manager,
api=True,
work_folder=options.output_folder,
)
with setup_pipeline(
options=options, plugin_manager=plugin_manager, api=True, work_folder=None
) as (work_folder, executor, plugin_manager):
origin_pdf = work_folder / 'origin.pdf'
shutil.copy2(options.input_file, origin_pdf)
origin_pdf = work_folder / 'origin.pdf'
shutil.copy2(options.input_file, origin_pdf)
# Gather pdfinfo and create context
pdfinfo = get_pdfinfo(
options.input_file,
executor=executor,
detailed_analysis=options.redo_ocr,
progbar=options.progress_bar,
max_workers=options.jobs if not options.use_threads else 1, # To help debug
check_pages=options.pages,
)
context = PdfContext(
options, work_folder, options.input_file, pdfinfo, plugin_manager
)
# Validate options are okay for this pdf
validate_pdfinfo_options(context)
optimize_messages = exec_hocr_to_ocr_pdf(context, executor)
# Gather pdfinfo and create context
pdfinfo = get_pdfinfo(
options.input_file,
executor=executor,
detailed_analysis=options.redo_ocr,
progbar=options.progress_bar,
max_workers=options.jobs if not options.use_threads else 1, # To help debug
check_pages=options.pages,
)
context = PdfContext(
options, work_folder, options.input_file, pdfinfo, plugin_manager
)
# Validate options are okay for this pdf
validate_pdfinfo_options(context)
optimize_messages = exec_hocr_to_ocr_pdf(context, executor)
report_output_pdf(options, origin_pdf, optimize_messages)
report_output_pdf(options, origin_pdf, optimize_messages)
src/ocrmypdf/_pipelines/pdf_to_hocr.py
+20 -20
View File
@@ -95,28 +95,28 @@ def run_hocr_pipeline(
*,
plugin_manager: OcrmypdfPluginManager | None,
) -> None:
work_folder, debug_log_handler, executor, plugin_manager = setup_pipeline(
options.keep_temporary_files = True
with setup_pipeline(
options=options,
plugin_manager=plugin_manager,
api=True,
work_folder=options.output_folder,
)
) as (work_folder, executor, plugin_manager):
shutil.copy2(options.input_file, work_folder / 'origin.pdf')
shutil.copy2(options.input_file, work_folder / 'origin.pdf')
# Gather pdfinfo and create context
pdfinfo = get_pdfinfo(
options.input_file,
executor=executor,
detailed_analysis=options.redo_ocr,
progbar=options.progress_bar,
max_workers=options.jobs if not options.use_threads else 1, # To help debug
check_pages=options.pages,
)
context = PdfContext(
options, work_folder, options.input_file, pdfinfo, plugin_manager
)
# Validate options are okay for this pdf
set_lossless_reconstruction(options)
validate_pdfinfo_options(context)
exec_pdf_to_hocr(context, executor)
# Gather pdfinfo and create context
pdfinfo = get_pdfinfo(
options.input_file,
executor=executor,
detailed_analysis=options.redo_ocr,
progbar=options.progress_bar,
max_workers=options.jobs if not options.use_threads else 1, # To help debug
check_pages=options.pages,
)
context = PdfContext(
options, work_folder, options.input_file, pdfinfo, plugin_manager
)
# Validate options are okay for this pdf
set_lossless_reconstruction(options)
validate_pdfinfo_options(context)
exec_pdf_to_hocr(context, executor)
src/ocrmypdf/_pipelines/standard.py
+66 -68
View File
@@ -179,81 +179,79 @@ def run_pipeline(
For CLI (``api=False``), exceptions are printed and described;
for API use, they are propagated to the caller.
"""
work_folder, debug_log_handler, executor, plugin_manager = setup_pipeline(
with setup_pipeline(
options=options, plugin_manager=plugin_manager, api=api, work_folder=None
)
try:
check_requested_output_file(options)
start_input_file, original_filename = create_input_file(options, work_folder)
) as (work_folder, executor, plugin_manager):
try:
check_requested_output_file(options)
start_input_file, original_filename = create_input_file(
options, work_folder
)
# Triage image or pdf
origin_pdf = triage(
original_filename, start_input_file, work_folder / 'origin.pdf', options
)
# Triage image or pdf
origin_pdf = triage(
original_filename, start_input_file, work_folder / 'origin.pdf', options
)
# Gather pdfinfo and create context
pdfinfo = get_pdfinfo(
origin_pdf,
executor=executor,
detailed_analysis=options.redo_ocr,
progbar=options.progress_bar,
max_workers=options.jobs if not options.use_threads else 1, # To help debug
check_pages=options.pages,
)
# Gather pdfinfo and create context
pdfinfo = get_pdfinfo(
origin_pdf,
executor=executor,
detailed_analysis=options.redo_ocr,
progbar=options.progress_bar,
max_workers=options.jobs
if not options.use_threads
else 1, # To help debug
check_pages=options.pages,
)
context = PdfContext(options, work_folder, origin_pdf, pdfinfo, plugin_manager)
context = PdfContext(
options, work_folder, origin_pdf, pdfinfo, plugin_manager
)
# Validate options are okay for this pdf
validate_pdfinfo_options(context)
# Validate options are okay for this pdf
validate_pdfinfo_options(context)
# Execute the pipeline
optimize_messages = exec_concurrent(context, executor)
# Execute the pipeline
optimize_messages = exec_concurrent(context, executor)
report_output_pdf(options, start_input_file, optimize_messages)
report_output_pdf(options, start_input_file, optimize_messages)
except KeyboardInterrupt if not api else NeverRaise:
if options.verbose >= 1:
log.exception("KeyboardInterrupt")
else:
log.error("KeyboardInterrupt")
return ExitCode.ctrl_c
except ExitCodeException if not api else NeverRaise as e:
e = cast(ExitCodeException, e)
if options.verbose >= 1:
log.exception("ExitCodeException")
elif str(e):
log.error("%s: %s", type(e).__name__, str(e))
else:
log.error(type(e).__name__)
return e.exit_code
except PIL.Image.DecompressionBombError if not api else NeverRaise:
log.exception(
"A decompression bomb error was encountered while executing the "
"pipeline. Use the argument --max-image-mpixels to raise the maximum "
"image pixel limit."
)
return ExitCode.other_error
except (
BrokenProcessPool if not api else NeverRaise,
BrokenThreadPool if not api else NeverRaise,
):
log.exception(
"A worker process was terminated unexpectedly. This is known to occur if "
"processing your file takes all available swap space and RAM. It may "
"help to try again with a smaller number of jobs, using the --jobs "
"argument."
)
return ExitCode.child_process_error
except Exception if not api else NeverRaise: # pylint: disable=broad-except
log.exception("An exception occurred while executing the pipeline")
return ExitCode.other_error
finally:
if debug_log_handler:
try:
debug_log_handler.close()
log.removeHandler(debug_log_handler)
except OSError as e:
print(e, file=sys.stderr)
cleanup_working_files(work_folder, options)
except KeyboardInterrupt if not api else NeverRaise:
if options.verbose >= 1:
log.exception("KeyboardInterrupt")
else:
log.error("KeyboardInterrupt")
return ExitCode.ctrl_c
except ExitCodeException if not api else NeverRaise as e:
e = cast(ExitCodeException, e)
if options.verbose >= 1:
log.exception("ExitCodeException")
elif str(e):
log.error("%s: %s", type(e).__name__, str(e))
else:
log.error(type(e).__name__)
return e.exit_code
except PIL.Image.DecompressionBombError if not api else NeverRaise:
log.exception(
"A decompression bomb error was encountered while executing the "
"pipeline. Use the argument --max-image-mpixels to raise the maximum "
"image pixel limit."
)
return ExitCode.other_error
except (
BrokenProcessPool if not api else NeverRaise,
BrokenThreadPool if not api else NeverRaise,
):
log.exception(
"A worker process was terminated unexpectedly. This is known to occur if "
"processing your file takes all available swap space and RAM. It may "
"help to try again with a smaller number of jobs, using the --jobs "
"argument."
)
return ExitCode.child_process_error
except Exception if not api else NeverRaise: # pylint: disable=broad-except
log.exception("An exception occurred while executing the pipeline")
return ExitCode.other_error
return ExitCode.ok