From 46a279a49aa4bd42896bcde842ae36691bd00104 Mon Sep 17 00:00:00 2001 From: "James R. Barlow" Date: Thu, 19 Oct 2023 00:20:28 -0700 Subject: [PATCH] Improve passing of arguments to workers The executor system was built around passing only a single argument to workers, which was always PageContext. For other tasks, all actual arguments were packed in a tuple, which meant we needed intermediate functions to unpack the tuple. The situation is now rationlized and resembles how Python handles argument passing to familiar multiprocessing tools. --- src/ocrmypdf/_concurrent.py | 2 +- src/ocrmypdf/_exec/jbig2enc.py | 16 ++-------------- src/ocrmypdf/_exec/pngquant.py | 4 ---- src/ocrmypdf/_jobcontext.py | 6 ++++++ src/ocrmypdf/_pipelines/hocr_to_ocr_pdf.py | 2 +- src/ocrmypdf/_pipelines/ocr.py | 2 +- src/ocrmypdf/_pipelines/pdf_to_hocr.py | 3 +-- src/ocrmypdf/builtin_plugins/concurrency.py | 2 +- src/ocrmypdf/extra_plugins/semfree.py | 4 ++-- src/ocrmypdf/optimize.py | 17 +++++++++-------- src/ocrmypdf/pdfinfo/info.py | 3 +-- 11 files changed, 25 insertions(+), 36 deletions(-) diff --git a/src/ocrmypdf/_concurrent.py b/src/ocrmypdf/_concurrent.py index 7d188365..9e10a16e 100644 --- a/src/ocrmypdf/_concurrent.py +++ b/src/ocrmypdf/_concurrent.py @@ -133,5 +133,5 @@ class SerialExecutor(Executor): ): # pylint: disable=unused-argument with self.pbar_class(**progress_kwargs) as pbar: for args in task_arguments: - result = task(args) + result = task(*args) task_finished(result, pbar) diff --git a/src/ocrmypdf/_exec/jbig2enc.py b/src/ocrmypdf/_exec/jbig2enc.py index 28d3dd1c..5a34a95a 100644 --- a/src/ocrmypdf/_exec/jbig2enc.py +++ b/src/ocrmypdf/_exec/jbig2enc.py @@ -25,7 +25,7 @@ def available(): return True -def convert_group(*, cwd, infiles, out_prefix, threshold): +def convert_group(cwd, infiles, out_prefix, threshold): args = [ 'jbig2', '-b', @@ -43,21 +43,9 @@ def convert_group(*, cwd, infiles, out_prefix, threshold): return proc -def convert_group_mp(args): - return convert_group( - cwd=args[0], infiles=args[1], out_prefix=args[2], threshold=args[3] - ) - - -def convert_single(*, cwd, infile, outfile, threshold): +def convert_single(cwd, infile, outfile, threshold): args = ['jbig2', '--pdf', '-t', str(threshold), infile] with open(outfile, 'wb') as fstdout: proc = run(args, cwd=cwd, stdout=fstdout, stderr=PIPE) proc.check_returncode() return proc - - -def convert_single_mp(args): - return convert_single( - cwd=args[0], infile=args[1], outfile=args[2], threshold=args[3] - ) diff --git a/src/ocrmypdf/_exec/pngquant.py b/src/ocrmypdf/_exec/pngquant.py index 3fa94277..5b8600d0 100644 --- a/src/ocrmypdf/_exec/pngquant.py +++ b/src/ocrmypdf/_exec/pngquant.py @@ -50,7 +50,3 @@ def quantize(input_file: Path, output_file: Path, quality_min: int, quality_max: if result.returncode == 0: # input_file could be the same as output_file, so we defer the write output_file.write_bytes(result.stdout) - - -def quantize_mp(args): - return quantize(*args) diff --git a/src/ocrmypdf/_jobcontext.py b/src/ocrmypdf/_jobcontext.py index 2ffaa353..0bc279a4 100644 --- a/src/ocrmypdf/_jobcontext.py +++ b/src/ocrmypdf/_jobcontext.py @@ -53,6 +53,12 @@ class PdfContext: for n in range(npages): yield PageContext(self, n) + def get_page_context_args(self) -> Iterator[tuple[PageContext]]: + """Get all ``PageContext`` for this PDF packaged in tuple for args-splatting.""" + npages = len(self.pdfinfo) + for n in range(npages): + yield (PageContext(self, n),) + class PageContext: """Holds our context for a page. diff --git a/src/ocrmypdf/_pipelines/hocr_to_ocr_pdf.py b/src/ocrmypdf/_pipelines/hocr_to_ocr_pdf.py index ea071eb6..200cc366 100644 --- a/src/ocrmypdf/_pipelines/hocr_to_ocr_pdf.py +++ b/src/ocrmypdf/_pipelines/hocr_to_ocr_pdf.py @@ -87,7 +87,7 @@ def exec_hocr_to_ocr_pdf(context: PdfContext, executor: Executor) -> Sequence[st ), worker_initializer=partial(worker_init, PIL.Image.MAX_IMAGE_PIXELS), task=exec_hocrtransform_sync, - task_arguments=context.get_page_contexts(), + task_arguments=context.get_page_context_args(), task_finished=graft_page, ) diff --git a/src/ocrmypdf/_pipelines/ocr.py b/src/ocrmypdf/_pipelines/ocr.py index 0a4bfd10..dacce069 100644 --- a/src/ocrmypdf/_pipelines/ocr.py +++ b/src/ocrmypdf/_pipelines/ocr.py @@ -128,7 +128,7 @@ def exec_concurrent(context: PdfContext, executor: Executor) -> Sequence[str]: ), worker_initializer=partial(worker_init, PIL.Image.MAX_IMAGE_PIXELS), task=exec_page_sync, - task_arguments=context.get_page_contexts(), + task_arguments=context.get_page_context_args(), task_finished=update_page, ) diff --git a/src/ocrmypdf/_pipelines/pdf_to_hocr.py b/src/ocrmypdf/_pipelines/pdf_to_hocr.py index d1a507bf..891878e6 100644 --- a/src/ocrmypdf/_pipelines/pdf_to_hocr.py +++ b/src/ocrmypdf/_pipelines/pdf_to_hocr.py @@ -39,7 +39,6 @@ from ocrmypdf._validation import ( log = logging.getLogger(__name__) - def exec_page_hocr_sync(page_context: PageContext) -> HOCRResult: """Execute a pipeline for a single page hOCR.""" set_thread_pageno(page_context.pageno + 1) @@ -82,7 +81,7 @@ def exec_pdf_to_hocr(context: PdfContext, executor: Executor) -> None: ), worker_initializer=partial(worker_init, PIL.Image.MAX_IMAGE_PIXELS), task=exec_page_hocr_sync, - task_arguments=context.get_page_contexts(), + task_arguments=context.get_page_context_args(), ) diff --git a/src/ocrmypdf/builtin_plugins/concurrency.py b/src/ocrmypdf/builtin_plugins/concurrency.py index d7080641..e2b80a59 100644 --- a/src/ocrmypdf/builtin_plugins/concurrency.py +++ b/src/ocrmypdf/builtin_plugins/concurrency.py @@ -134,7 +134,7 @@ class StandardExecutor(Executor): initializer=initializer, initargs=(log_queue, worker_initializer, logging.getLogger("").level), ) as executor: - futures = [executor.submit(task, args) for args in task_arguments] + futures = [executor.submit(task, *args) for args in task_arguments] try: for future in as_completed(futures): result = future.result() diff --git a/src/ocrmypdf/extra_plugins/semfree.py b/src/ocrmypdf/extra_plugins/semfree.py index 571d7257..6aeec84c 100644 --- a/src/ocrmypdf/extra_plugins/semfree.py +++ b/src/ocrmypdf/extra_plugins/semfree.py @@ -94,7 +94,7 @@ def process_loop( for args in task_args: try: - result = task(args) + result = task(*args) except Exception as e: # pylint: disable=broad-except conn.send((MessageType.exception, e)) break @@ -123,7 +123,7 @@ class LambdaExecutor(Executor): if use_threads and max_workers == 1: with self.pbar_class(**progress_kwargs) as pbar: for args in task_arguments: - result = task(args) + result = task(*args) task_finished(result, pbar) return diff --git a/src/ocrmypdf/optimize.py b/src/ocrmypdf/optimize.py index 13d30efe..9bdd54b1 100644 --- a/src/ocrmypdf/optimize.py +++ b/src/ocrmypdf/optimize.py @@ -390,10 +390,10 @@ def _produce_jbig2_images( if options.jbig2_page_group_size > 1: jbig2_args = jbig2_group_args - jbig2_convert = jbig2enc.convert_group_mp + jbig2_convert = jbig2enc.convert_group else: jbig2_args = jbig2_single_args - jbig2_convert = jbig2enc.convert_single_mp + jbig2_convert = jbig2enc.convert_single executor( use_threads=True, @@ -454,9 +454,9 @@ def convert_to_jbig2( ) -def _optimize_jpeg(args: tuple[Xref, Path, Path, int]) -> tuple[Xref, Path | None]: - xref, in_jpg, opt_jpg, jpeg_quality = args - +def _optimize_jpeg( + xref: Xref, in_jpg: Path, opt_jpg: Path, jpeg_quality: int +) -> tuple[Xref, Path | None]: with Image.open(in_jpg) as im: im.save(opt_jpg, optimize=True, quality=jpeg_quality) @@ -515,8 +515,9 @@ def _find_deflatable_jpeg( return None -def _deflate_jpeg(args: tuple[Pdf, threading.Lock, Xref, int]) -> tuple[Xref, bytes]: - pdf, lock, xref, complevel = args +def _deflate_jpeg( + pdf: Pdf, lock: threading.Lock, xref: Xref, complevel: int +) -> tuple[Xref, bytes]: with lock: xobj = pdf.get_object(xref, 0) try: @@ -651,7 +652,7 @@ def transcode_pngs( unit='image', disable=not options.progress_bar, ), - task=pngquant.quantize_mp, + task=pngquant.quantize, task_arguments=pngquant_args(), ) diff --git a/src/ocrmypdf/pdfinfo/info.py b/src/ocrmypdf/pdfinfo/info.py index eb0bf350..7815ad21 100644 --- a/src/ocrmypdf/pdfinfo/info.py +++ b/src/ocrmypdf/pdfinfo/info.py @@ -694,8 +694,7 @@ def _pdf_pageinfo_sync_init(pdf: Pdf, infile: Path, pdfminer_loglevel): atexit.register(on_process_close) -def _pdf_pageinfo_sync(args): - pageno, thread_pdf, infile, check_pages, detailed_analysis = args +def _pdf_pageinfo_sync(pageno, thread_pdf, infile, check_pages, detailed_analysis): pdf = thread_pdf if thread_pdf is not None else worker_pdf with ExitStack() as stack: if not pdf: # When called with SerialExecutor