Refactor setup_pipeline to decouple manage_work_folder

This commit is contained in:
James R. Barlow
2023-10-24 00:54:30 -07:00
parent fc6f959d21
commit 6f82097d14
4 changed files with 37 additions and 22 deletions
+4 -12
View File
@@ -187,9 +187,7 @@ def manage_debug_log_handler(
@contextmanager
def manage_work_folder(*, work_folder: Path | None, retain: bool, print_location: bool):
if not work_folder:
work_folder = Path(mkdtemp(prefix="ocrmypdf.io."))
def manage_work_folder(*, work_folder: Path, retain: bool, print_location: bool):
try:
yield work_folder
finally:
@@ -210,7 +208,7 @@ def setup_pipeline(
plugin_manager: OcrmypdfPluginManager | None,
api: bool = False,
work_folder: Path,
) -> Generator[tuple[Path, Executor, OcrmypdfPluginManager], None, None]:
) -> Generator[tuple[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.)
@@ -219,16 +217,10 @@ def setup_pipeline(
if not plugin_manager:
plugin_manager = get_plugin_manager(options.plugins)
with manage_work_folder(
work_folder=work_folder,
retain=options.keep_temporary_files,
print_location=True,
) as work_folder, manage_debug_log_handler(
options=options, api=api, work_folder=work_folder
):
with manage_debug_log_handler(options=options, api=api, work_folder=work_folder):
pikepdf_enable_mmap()
executor = setup_executor(plugin_manager)
yield work_folder, executor, plugin_manager
yield executor, plugin_manager
def preprocess(
+12 -3
View File
@@ -28,6 +28,7 @@ from ocrmypdf._pipeline import (
)
from ocrmypdf._pipelines.common import (
HOCRResult,
manage_work_folder,
post_process,
report_output_pdf,
set_logging_tls,
@@ -112,9 +113,17 @@ def run_hocr_to_ocr_pdf_pipeline(
*,
plugin_manager: OcrmypdfPluginManager | None,
) -> ExitCode:
with setup_pipeline(
options=options, plugin_manager=plugin_manager, api=True, work_folder=None
) as (work_folder, executor, plugin_manager):
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,
):
origin_pdf = work_folder / 'origin.pdf'
shutil.copy2(options.input_file, origin_pdf)
+9 -4
View File
@@ -26,6 +26,7 @@ from ocrmypdf._pipeline import (
)
from ocrmypdf._pipelines.common import (
HOCRResult,
manage_work_folder,
process_page,
set_logging_tls,
setup_pipeline,
@@ -95,13 +96,17 @@ def run_hocr_pipeline(
*,
plugin_manager: OcrmypdfPluginManager | None,
) -> None:
options.keep_temporary_files = True
with setup_pipeline(
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=options.output_folder,
) as (work_folder, executor, plugin_manager):
work_folder=work_folder,
) as (
executor,
plugin_manager,
):
shutil.copy2(options.input_file, work_folder / 'origin.pdf')
# Gather pdfinfo and create context
+12 -3
View File
@@ -17,6 +17,7 @@ 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
@@ -37,6 +38,7 @@ from ocrmypdf._pipeline import (
)
from ocrmypdf._pipelines.common import (
PageResult,
manage_work_folder,
post_process,
process_page,
report_output_pdf,
@@ -179,9 +181,16 @@ def run_pipeline(
For CLI (``api=False``), exceptions are printed and described;
for API use, they are propagated to the caller.
"""
with setup_pipeline(
options=options, plugin_manager=plugin_manager, api=api, work_folder=None
) as (work_folder, executor, plugin_manager):
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(