From 6f82097d141cb7d2ae76487aec654e7787b2a89a Mon Sep 17 00:00:00 2001
From: "James R. Barlow" <james@purplerock.ca>
Date: Thu, 12 Oct 2023 23:43:49 -0700
Subject: [PATCH] Refactor setup_pipeline to decouple manage_work_folder

---
 src/ocrmypdf/_pipelines/common.py          | 16 ++++------------
 src/ocrmypdf/_pipelines/hocr_to_ocr_pdf.py | 15 ++++++++++++---
 src/ocrmypdf/_pipelines/pdf_to_hocr.py     | 13 +++++++++----
 src/ocrmypdf/_pipelines/standard.py        | 15 ++++++++++++---
 4 files changed, 37 insertions(+), 22 deletions(-)

diff --git a/src/ocrmypdf/_pipelines/common.py b/src/ocrmypdf/_pipelines/common.py
index 05e4ccee..5ba73479 100644
--- a/src/ocrmypdf/_pipelines/common.py
+++ b/src/ocrmypdf/_pipelines/common.py
@@ -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(
diff --git a/src/ocrmypdf/_pipelines/hocr_to_ocr_pdf.py b/src/ocrmypdf/_pipelines/hocr_to_ocr_pdf.py
index 44623292..7a4d3859 100644
--- a/src/ocrmypdf/_pipelines/hocr_to_ocr_pdf.py
+++ b/src/ocrmypdf/_pipelines/hocr_to_ocr_pdf.py
@@ -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)
 
diff --git a/src/ocrmypdf/_pipelines/pdf_to_hocr.py b/src/ocrmypdf/_pipelines/pdf_to_hocr.py
index fdc74958..1cfb9851 100644
--- a/src/ocrmypdf/_pipelines/pdf_to_hocr.py
+++ b/src/ocrmypdf/_pipelines/pdf_to_hocr.py
@@ -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
diff --git a/src/ocrmypdf/_pipelines/standard.py b/src/ocrmypdf/_pipelines/standard.py
index 4a78fb3e..b2dd36f2 100644
--- a/src/ocrmypdf/_pipelines/standard.py
+++ b/src/ocrmypdf/_pipelines/standard.py
@@ -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(