Refactor setup_pipeline to decouple manage_work_folder
This commit is contained in:
@@ -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(
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user