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

Refactor CLI exception handling

Browse Source
This commit is contained in:
James R. Barlow
2023-10-24 00:54:30 -07:00
parent 6f82097d14
commit 67be4d1904
4 changed files with 116 additions and 131 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/ocrmypdf/_pipelines/common.py
+58 -19
View File
@@ -1,3 +1,6 @@
# SPDX-FileCopyrightText: 2023 James R. Barlow
# SPDX-License-Identifier: MPL-2.0
from __future__ import annotations
import argparse
@@ -7,12 +10,13 @@ import logging.handlers
import os
import shutil
import sys
from collections.abc import Generator, Sequence
from collections.abc import Sequence
from concurrent.futures.process import BrokenProcessPool
from concurrent.futures.thread import BrokenThreadPool
from contextlib import contextmanager
from dataclasses import dataclass
from pathlib import Path
from tempfile import mkdtemp
from typing import NamedTuple
from typing import Any, Callable, NamedTuple, cast
import PIL
@@ -39,7 +43,7 @@ from ocrmypdf._plugin_manager import OcrmypdfPluginManager, get_plugin_manager
from ocrmypdf._validation import (
report_output_file_size,
)
from ocrmypdf.exceptions import ExitCode
from ocrmypdf.exceptions import ExitCode, ExitCodeException
from ocrmypdf.helpers import (
available_cpu_count,
check_pdf,
@@ -160,14 +164,11 @@ def worker_init(max_pixels: int) -> None:
def manage_debug_log_handler(
*,
options: argparse.Namespace,
api: bool = False,
work_folder: Path | None,
work_folder: Path,
):
debug_log_handler = None
if (
(options.keep_temporary_files or options.verbose >= 1)
and not os.environ.get('PYTEST_CURRENT_TEST', '')
and not api
if (options.keep_temporary_files or options.verbose >= 1) and not os.environ.get(
'PYTEST_CURRENT_TEST', ''
):
# Debug log for command line interface only with verbose output
# See https://github.com/pytest-dev/pytest/issues/5502 for why we skip this
@@ -201,14 +202,53 @@ def manage_work_folder(*, work_folder: Path, retain: bool, print_location: bool)
shutil.rmtree(work_folder, ignore_errors=True)
@contextmanager
def cli_exception_handler(
fn: Callable[[Any], ExitCode], *args, options: argparse.Namespace
) -> ExitCode:
try:
return fn(*args)
except KeyboardInterrupt:
if options.verbose >= 1:
log.exception("KeyboardInterrupt")
else:
log.error("KeyboardInterrupt")
return ExitCode.ctrl_c
except ExitCodeException 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:
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,
BrokenThreadPool,
):
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: # pylint: disable=broad-except
log.exception("An exception occurred while executing the pipeline")
return ExitCode.other_error
def setup_pipeline(
*,
options: argparse.Namespace,
plugin_manager: OcrmypdfPluginManager | None,
api: bool = False,
work_folder: Path,
) -> Generator[tuple[Executor, OcrmypdfPluginManager], None, None]:
) -> tuple[Executor, OcrmypdfPluginManager]:
# 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.)
@@ -217,10 +257,9 @@ def setup_pipeline(
if not plugin_manager:
plugin_manager = get_plugin_manager(options.plugins)
with manage_debug_log_handler(options=options, api=api, work_folder=work_folder):
pikepdf_enable_mmap()
executor = setup_executor(plugin_manager)
yield executor, plugin_manager
pikepdf_enable_mmap()
executor = setup_executor(plugin_manager)
return executor, plugin_manager
def preprocess(
src/ocrmypdf/_pipelines/hocr_to_ocr_pdf.py
+3 -10
View File
@@ -1,4 +1,4 @@
# SPDX-FileCopyrightText: 2019-2022 James R. Barlow
# SPDX-FileCopyrightText: 2019-2023 James R. Barlow
# SPDX-FileCopyrightText: 2019 Martin Wind
# SPDX-License-Identifier: MPL-2.0
@@ -115,15 +115,8 @@ def run_hocr_to_ocr_pdf_pipeline(
) -> ExitCode:
with manage_work_folder(
work_folder=options.input_folder, retain=True, print_location=False
) as work_folder, setup_pipeline(
options=options,
plugin_manager=plugin_manager,
api=True,
work_folder=work_folder,
) as (
executor,
plugin_manager,
):
) as work_folder:
executor, plugin_manager = setup_pipeline(options, plugin_manager)
origin_pdf = work_folder / 'origin.pdf'
shutil.copy2(options.input_file, origin_pdf)
src/ocrmypdf/_pipelines/pdf_to_hocr.py
+3 -10
View File
@@ -1,4 +1,4 @@
# SPDX-FileCopyrightText: 2019-2022 James R. Barlow
# SPDX-FileCopyrightText: 2019-2023 James R. Barlow
# SPDX-FileCopyrightText: 2019 Martin Wind
# SPDX-License-Identifier: MPL-2.0
@@ -98,15 +98,8 @@ def run_hocr_pipeline(
) -> None:
with manage_work_folder(
work_folder=options.output_folder, retain=True, print_location=False
) as work_folder, setup_pipeline(
options=options,
plugin_manager=plugin_manager,
api=True,
work_folder=work_folder,
) as (
executor,
plugin_manager,
):
) as work_folder:
executor, plugin_manager = setup_pipeline(options, plugin_manager)
shutil.copy2(options.input_file, work_folder / 'origin.pdf')
# Gather pdfinfo and create context
src/ocrmypdf/_pipelines/standard.py
+52 -92
View File
@@ -1,4 +1,4 @@
# SPDX-FileCopyrightText: 2019-2022 James R. Barlow
# SPDX-FileCopyrightText: 2019-2023 James R. Barlow
# SPDX-FileCopyrightText: 2019 Martin Wind
# SPDX-License-Identifier: MPL-2.0
@@ -10,15 +10,11 @@ from __future__ import annotations
import argparse
import logging
import logging.handlers
import sys
import threading
from collections.abc import Sequence
from concurrent.futures.process import BrokenProcessPool
from concurrent.futures.thread import BrokenThreadPool
from functools import partial
from pathlib import Path
from tempfile import mkdtemp
from typing import cast
import PIL
@@ -38,6 +34,8 @@ from ocrmypdf._pipeline import (
)
from ocrmypdf._pipelines.common import (
PageResult,
cli_exception_handler,
manage_debug_log_handler,
manage_work_folder,
post_process,
process_page,
@@ -51,10 +49,7 @@ from ocrmypdf._validation import (
check_requested_output_file,
create_input_file,
)
from ocrmypdf.exceptions import ExitCode, ExitCodeException
from ocrmypdf.helpers import (
NeverRaise,
)
from ocrmypdf.exceptions import ExitCode
log = logging.getLogger(__name__)
@@ -164,6 +159,48 @@ def exec_concurrent(context: PdfContext, executor: Executor) -> Sequence[str]:
return messages
def _run_pipeline(
options: argparse.Namespace,
plugin_manager: OcrmypdfPluginManager | None,
) -> ExitCode:
with manage_work_folder(
work_folder=Path(mkdtemp(prefix="ocrmypdf.io.")),
retain=options.keep_temporary_files,
print_location=options.keep_temporary_files,
) as work_folder, manage_debug_log_handler(
options=options, work_folder=work_folder
):
executor, plugin_manager = setup_pipeline(options, plugin_manager)
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
)
# 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)
# Validate options are okay for this pdf
validate_pdfinfo_options(context)
# Execute the pipeline
optimize_messages = exec_concurrent(context, executor)
exitcode = report_output_pdf(options, start_input_file, optimize_messages)
return exitcode
def run_pipeline(
options: argparse.Namespace,
*,
@@ -181,86 +218,9 @@ def run_pipeline(
For CLI (``api=False``), exceptions are printed and described;
for API use, they are propagated to the caller.
"""
with manage_work_folder(
work_folder=Path(mkdtemp(prefix="ocrmypdf.io.")),
retain=options.keep_temporary_files,
print_location=options.keep_temporary_files,
) as work_folder, setup_pipeline(
options=options, plugin_manager=plugin_manager, api=api, work_folder=work_folder
) as (
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
)
# 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
)
# Validate options are okay for this pdf
validate_pdfinfo_options(context)
# Execute the pipeline
optimize_messages = exec_concurrent(context, executor)
exitcode = 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
return exitcode
if api:
return _run_pipeline(options, plugin_manager)
else:
return cli_exception_handler(
_run_pipeline, options, plugin_manager, options=options
)