Support input/output streams at API level

src/ocrmypdf/_jobcontext.py

@@ -19,6 +19,8 @@ import os
 import shutil
 import sys
 from argparse import Namespace
+from copy import copy
+from io import IOBase
 from pathlib import Path
 from typing import Iterator
 
@@ -41,9 +43,6 @@ class PdfContext:
         self.origin = origin
         self.pdfinfo = pdfinfo
         self.plugin_manager = plugin_manager
-        self.name = os.path.basename(options.input_file)
-        if self.name == '-':
-            self.name = 'stdin'
 
     def get_path(self, name: str) -> Path:
         return self.work_folder / name
@@ -65,7 +64,6 @@ class PageContext:
         self.work_folder = pdf_context.work_folder
         self.origin = pdf_context.origin
         self.options = pdf_context.options
-        self.name = pdf_context.name
         self.pageno = pageno
         self.pageinfo = pdf_context.pdfinfo[pageno]
         self.plugin_manager = pdf_context.plugin_manager
@@ -73,6 +71,16 @@ class PageContext:
     def get_path(self, name: str) -> Path:
         return self.work_folder / ("%06d_%s" % (self.pageno + 1, name))
 
+    def __getstate__(self):
+        state = self.__dict__.copy()
+
+        state['options'] = copy(self.options)
+        if not isinstance(state['options'].input_file, (str, bytes, os.PathLike)):
+            state['options'].input_file = 'stream'
+        if not isinstance(state['options'].output_file, (str, bytes, os.PathLike)):
+            state['options'].output_file = 'stream'
+        return state
+
 
 def cleanup_working_files(work_folder: Path, options: Namespace):
     if options.keep_temporary_files:

src/ocrmypdf/_pipeline.py

@@ -19,6 +19,7 @@ import logging
 import os
 import re
 import sys
+from contextlib import suppress
 from datetime import datetime, timezone
 from pathlib import Path
 from shutil import copyfileobj
@@ -822,6 +823,10 @@ def copy_final(input_file: Path, output_file: Path, _context: PdfContext):
         if output_file == '-':
             copyfileobj(input_stream, sys.stdout.buffer)
             sys.stdout.flush()
+        elif hasattr(output_file, 'writable'):
+            copyfileobj(input_stream, output_file)
+            with suppress(AttributeError):
+                output_file.flush()
         else:
             # At this point we overwrite the output_file specified by the user
             # use copyfileobj because then we use open() to create the file and

src/ocrmypdf/_sync.py

@@ -357,6 +357,10 @@ def run_pipeline(options, *, plugin_manager, api=False):
 
         if options.output_file == '-':
             log.info("Output sent to stdout")
+        elif (
+            hasattr(options.output_file, 'writable') and options.output_file.writable()
+        ):
+            log.info("Output written to stream")
         elif samefile(options.output_file, os.devnull):
             pass  # Say nothing when sending to dev null
         else:

src/ocrmypdf/_validation.py

@@ -337,7 +337,15 @@ def create_input_file(options, work_folder: Path) -> Tuple[Path, str]:
         target = work_folder / 'stdin'
         with open(target, 'wb') as stream_buffer:
             copyfileobj(sys.stdin.buffer, stream_buffer)
-        return target, "<stdin>"
+        return target, "stdin"
+    elif hasattr(options.input_file, 'readable'):
+        if not options.input_file.readable():
+            raise InputFileError("Input file stream is not readable")
+        log.info('reading file from input stream')
+        target = os.path.join(work_folder, 'stream')
+        with open(target, 'wb') as stream_buffer:
+            copyfileobj(options.input_file, stream_buffer)
+        return target, "stream"
     else:
         try:
             target = work_folder / 'origin'
@@ -355,6 +363,9 @@ def check_requested_output_file(options):
                 "is connected to a terminal.  Please redirect stdout to a "
                 "file."
             )
+    elif hasattr(options.output_file, 'writable'):
+        if not options.output_file.writable():
+            raise OutputFileAccessError("Output stream is not writable")
     elif not is_file_writable(options.output_file):
         raise OutputFileAccessError(
             f"Output file location ({options.output_file}) is not a writable file."

src/ocrmypdf/api.py

@@ -21,7 +21,7 @@ import sys
 from argparse import ArgumentParser
 from enum import IntEnum
 from pathlib import Path
-from typing import Iterable
+from typing import BinaryIO, Iterable, Union
 
 from ocrmypdf._logging import PageNumberFilter, TqdmConsole
 from ocrmypdf._plugin_manager import get_plugin_manager
@@ -35,6 +35,9 @@ except ModuleNotFoundError:
     coloredlogs = None
 
 
+PathOrIO = Union[BinaryIO, os.PathLike]
+
+
 class Verbosity(IntEnum):
     """Verbosity level for configure_logging."""
 
@@ -127,11 +130,7 @@ def configure_logging(
 
 
 def create_options(
-    *,
-    input_file: os.PathLike,
-    output_file: os.PathLike,
-    parser: ArgumentParser,
-    **kwargs,
+    *, input_file: PathOrIO, output_file: PathOrIO, parser: ArgumentParser, **kwargs
 ):
     cmdline = []
     deferred = []
@@ -171,19 +170,31 @@ def create_options(
         else:
             raise TypeError(f"{arg}: {val} ({type(val)})")
 
-    cmdline.append(str(input_file))
-    cmdline.append(str(output_file))
+    try:
+        cmdline.append(os.fspath(input_file))
+    except TypeError:
+        cmdline.append('stream://input_file')
+    try:
+        cmdline.append(os.fspath(output_file))
+    except TypeError:
+        cmdline.append('stream://output_file')
 
     parser._api_mode = True
     options = parser.parse_args(cmdline)
     for keyword, val in deferred:
         setattr(options, keyword, val)
+
+    if options.input_file == 'stream://input_file':
+        options.input_file = input_file
+    if options.output_file == 'stream://output_file':
+        options.output_file = output_file
+
     return options
 
 
 def ocr(  # pylint: disable=unused-argument
-    input_file: os.PathLike,
-    output_file: os.PathLike,
+    input_file: PathOrIO,
+    output_file: PathOrIO,
     *,
     language: Iterable[str] = None,
     image_dpi: int = None,

tests/test_api.py

@@ -16,7 +16,7 @@
 # along with OCRmyPDF.  If not, see <http://www.gnu.org/licenses/>.
 
 import logging
-from io import StringIO
+from io import BytesIO, StringIO
 
 import pytest
 from tqdm import tqdm
@@ -66,3 +66,12 @@ def test_language_list():
         (ocrmypdf.exceptions.InputFileError, ocrmypdf.exceptions.MissingDependencyError)
     ):
         ocrmypdf.ocr('doesnotexist.pdf', '_.pdf', language=['eng', 'deu'])
+
+
+def test_stream_api(resources):
+    in_ = (resources / 'graph.pdf').open('rb')
+    out = BytesIO()
+
+    ocrmypdf.ocr(in_, out, tesseract_timeout=0.0)
+    out.seek(0)
+    assert b'%PDF' in out.read(1024)