refactor: post-AI code cleanup

This commit is contained in:
James R. Barlow
2025-12-21 12:21:47 -08:00
parent e1d976168c
commit 1f493ba789
19 changed files with 119 additions and 119 deletions
-1
View File
@@ -18,7 +18,6 @@ def main(
engine: Annotated[str, typer.Option()] = 'pdftotext',
):
"""Compare text in PDFs."""
text1 = run(
['pdftotext', '-layout', '-', '-'], stdin=pdf1, capture_output=True, check=True
)
-1
View File
@@ -9,7 +9,6 @@ import logging
import os
import re
from collections import deque
from io import BytesIO
from os import fspath
from pathlib import Path
from subprocess import PIPE, CalledProcessError
+7 -9
View File
@@ -5,12 +5,9 @@
from __future__ import annotations
import os
from collections.abc import Iterator
from copy import copy
from argparse import Namespace
from collections.abc import Iterator
from pathlib import Path
from typing import Union
from pluggy import PluginManager
@@ -29,7 +26,7 @@ class PdfContext:
def __init__(
self,
options: Union[OCROptions, Namespace],
options: OCROptions | Namespace,
work_folder: Path,
origin: Path,
pdfinfo: PdfInfo,
@@ -41,7 +38,7 @@ class PdfContext:
else:
# Convert Namespace to OCROptions
self.options = OCROptions.from_namespace(options)
self.work_folder = work_folder
self.origin = origin
self.pdfinfo = pdfinfo
@@ -73,7 +70,7 @@ class PageContext:
Must be pickle-able, so stores only intrinsic/simple data elements or those
capable of their serializing themselves via ``__getstate__``.
Note: Uses OCROptions with JSON serialization for multiprocessing compatibility.
"""
@@ -114,7 +111,6 @@ class PageContext:
# Fallback: if JSON serialization fails, convert to namespace
# This shouldn't happen but provides safety
from argparse import Namespace
import os
clean_options = Namespace()
for key, value in vars(self.options.to_namespace()).items():
@@ -122,6 +118,7 @@ class PageContext:
continue
try:
import pickle
pickle.dumps(value)
setattr(clean_options, key, value)
except TypeError:
@@ -134,10 +131,11 @@ class PageContext:
def __setstate__(self, state):
self.__dict__.update(state)
# Reconstruct OCROptions from JSON if available
if 'options_json' in state:
from ocrmypdf._options import OCROptions
self.options = OCROptions.model_validate_json_safe(state['options_json'])
# Otherwise, we have a fallback Namespace (shouldn't happen in normal operation)
# Leave it as-is for compatibility
-1
View File
@@ -15,7 +15,6 @@ from pikepdf import Dictionary, Name, Pdf
from pikepdf import __version__ as PIKEPDF_VERSION
from pikepdf.models.metadata import PdfMetadata, encode_pdf_date
from ocrmypdf._annots import remove_broken_goto_annotations
from ocrmypdf._defaults import PROGRAM_NAME
from ocrmypdf._jobcontext import PdfContext
from ocrmypdf._version import __version__ as OCRMYPF_VERSION
+43 -37
View File
@@ -10,11 +10,10 @@ import logging
import os
import unicodedata
from argparse import Namespace
from collections.abc import Iterable, Sequence
from copy import copy
from collections.abc import Sequence
from io import IOBase
from pathlib import Path
from typing import Any, BinaryIO, Union
from typing import Any, BinaryIO
from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
@@ -24,7 +23,7 @@ from ocrmypdf.helpers import monotonic
log = logging.getLogger(__name__)
PathOrIO = Union[BinaryIO, IOBase, Path, str, bytes]
PathOrIO = BinaryIO | IOBase | Path | str | bytes
def _pages_from_ranges(ranges: str) -> set[int]:
@@ -124,14 +123,14 @@ class OCROptions(BaseModel):
png_quality: int | None = None
jbig2_lossy: bool | None = None
jbig2_page_group_size: int | None = None
jbig2_threshold: float = 0.85
jbig2_threshold: float = 0.85
# Compatibility alias for plugins that expect jpeg_quality
@property
def jpeg_quality(self):
"""Compatibility alias for jpg_quality."""
return self.jpg_quality
@jpeg_quality.setter
def jpeg_quality(self, value):
"""Compatibility alias for jpg_quality."""
@@ -140,26 +139,25 @@ class OCROptions(BaseModel):
# Advanced options
max_image_mpixels: float = 250.0
pdf_renderer: str = 'auto'
tesseract_config: list[str] = []
tesseract_config: list[str] = []
tesseract_pagesegmode: int | None = None
tesseract_oem: int | None = None
tesseract_thresholding: int | None = None
tesseract_timeout: float = 0.0
tesseract_timeout: float = 0.0
tesseract_non_ocr_timeout: float | None = None
tesseract_downsample_above: int = 32767
tesseract_downsample_above: int = 32767
tesseract_downsample_large_images: bool | None = None
rotate_pages_threshold: float = DEFAULT_ROTATE_PAGES_THRESHOLD
pdfa_image_compression: str | None = None
color_conversion_strategy: str = "LeaveColorUnchanged"
color_conversion_strategy: str = "LeaveColorUnchanged"
user_words: os.PathLike | None = None
user_patterns: os.PathLike | None = None
fast_web_view: float = 1.0
fast_web_view: float = 1.0
continue_on_soft_render_error: bool | None = None
# Plugin system
plugins: Sequence[Path | str] | None = None
# Store any extra attributes (for plugins and dynamic options)
extra_attrs: dict[str, Any] = Field(
default_factory=dict, exclude=True, alias='_extra_attrs'
@@ -259,7 +257,6 @@ class OCROptions(BaseModel):
raise ValueError(f"pdf_renderer must be one of {valid_renderers}")
return v
@field_validator('clean_final')
@classmethod
def validate_clean_final(cls, v, info):
@@ -314,7 +311,7 @@ class OCROptions(BaseModel):
"""Validate metadata strings don't contain unsupported Unicode characters."""
if v is None:
return v
for char in v:
if unicodedata.category(char) == 'Co' or ord(char) >= 0x10000:
hexchar = hex(ord(char))[2:].upper()
@@ -332,7 +329,7 @@ class OCROptions(BaseModel):
return v
if isinstance(v, set):
return v # Already processed
# Convert string ranges to set of page numbers
return _pages_from_ranges(v)
@@ -356,10 +353,13 @@ class OCROptions(BaseModel):
raise ValueError("Choose only one of --force-ocr, --skip-text, --redo-ocr.")
return self
@model_validator(mode='after')
@model_validator(mode='after')
def validate_output_type_compatibility(self):
"""Validate output type is compatible with output file."""
if self.output_type == 'none' and str(self.output_file) not in (os.devnull, '-'):
if self.output_type == 'none' and str(self.output_file) not in (
os.devnull,
'-',
):
raise ValueError(
"Since you specified `--output-type none`, the output file "
f"{self.output_file} cannot be produced. Set the output file to "
@@ -370,19 +370,21 @@ class OCROptions(BaseModel):
@model_validator(mode='after')
def set_lossless_reconstruction(self):
"""Set lossless_reconstruction based on other options."""
lossless = not any([
self.deskew,
self.clean_final,
self.force_ocr,
self.remove_background,
])
lossless = not any(
[
self.deskew,
self.clean_final,
self.force_ocr,
self.remove_background,
]
)
if not lossless and self.redo_ocr:
raise ValueError(
"--redo-ocr is not currently compatible with --deskew, "
"--clean-final, and --remove-background"
)
# Set the computed attribute
self.extra_attrs['lossless_reconstruction'] = lossless
return self
@@ -391,12 +393,16 @@ class OCROptions(BaseModel):
"""Serialize to JSON with special handling for non-serializable types."""
# Create a copy of the model data for serialization
data = self.model_dump()
# Handle special types that don't serialize to JSON directly
def _serialize_value(value):
if isinstance(value, Path):
return {'__type__': 'Path', 'value': str(value)}
elif isinstance(value, (BinaryIO, IOBase)) or hasattr(value, 'read') or hasattr(value, 'write'):
elif (
isinstance(value, (BinaryIO, IOBase))
or hasattr(value, 'read')
or hasattr(value, 'write')
):
# Stream object - replace with placeholder
return {'__type__': 'Stream', 'value': 'stream'}
elif hasattr(value, '__class__') and 'Iterator' in value.__class__.__name__:
@@ -408,23 +414,23 @@ class OCROptions(BaseModel):
return {k: _serialize_value(v) for k, v in value.items()}
else:
return value
# Process all fields
serializable_data = {}
for key, value in data.items():
serializable_data[key] = _serialize_value(value)
# Add extra_attrs
if self.extra_attrs:
serializable_data['_extra_attrs'] = _serialize_value(self.extra_attrs)
return json.dumps(serializable_data)
@classmethod
def model_validate_json_safe(cls, json_str: str) -> OCROptions:
"""Reconstruct from JSON with special handling for non-serializable types."""
data = json.loads(json_str)
# Handle special types during deserialization
def _deserialize_value(value):
if isinstance(value, dict) and '__type__' in value:
@@ -441,21 +447,21 @@ class OCROptions(BaseModel):
return {k: _deserialize_value(v) for k, v in value.items()}
else:
return value
# Process all fields
deserialized_data = {}
extra_attrs = {}
for key, value in data.items():
if key == '_extra_attrs':
extra_attrs = _deserialize_value(value)
else:
deserialized_data[key] = _deserialize_value(value)
# Create instance
instance = cls(**deserialized_data)
instance.extra_attrs = extra_attrs
return instance
model_config = ConfigDict(
+4 -2
View File
@@ -915,10 +915,12 @@ def convert_to_pdfa(input_pdf: Path, input_ps_stub: Path, context: PdfContext) -
if options.output_type == 'pdfa':
pdfa_part = '2' # Default to PDF/A-2
else:
pdfa_part = options.output_type.split('-')[-1] # Extract number from pdfa-1, pdfa-2, etc.
pdfa_part = options.output_type.split('-')[
-1
] # Extract number from pdfa-1, pdfa-2, etc.
else:
pdfa_part = '2' # Fallback
context.plugin_manager.hook.generate_pdfa(
pdf_version=input_pdfinfo.min_version,
pdf_pages=[fix_docinfo_file],
+4 -6
View File
@@ -50,9 +50,8 @@ from ocrmypdf._plugin_manager import OcrmypdfPluginManager
from ocrmypdf._validation import (
report_output_file_size,
)
from ocrmypdf.exceptions import BadArgsError, ExitCode, ExitCodeException
from ocrmypdf.exceptions import ExitCode, ExitCodeException
from ocrmypdf.helpers import (
available_cpu_count,
check_pdf,
pikepdf_enable_mmap,
running_in_docker,
@@ -328,15 +327,14 @@ def setup_pipeline(
return executor
def do_get_pdfinfo(
pdf_path: Path, executor: Executor, options
) -> PdfInfo:
def do_get_pdfinfo(pdf_path: Path, executor: Executor, options) -> PdfInfo:
# Handle pages field - it might be a string that needs conversion
check_pages = options.pages
if isinstance(check_pages, str):
from ocrmypdf._options import _pages_from_ranges
check_pages = _pages_from_ranges(check_pages)
return get_pdfinfo(
pdf_path,
executor=executor,
+1 -1
View File
@@ -63,7 +63,7 @@ def _image_to_ocr_text(
pdf_renderer = options.pdf_renderer
if pdf_renderer == 'auto':
pdf_renderer = 'hocr'
if pdf_renderer.startswith('hocr'):
hocr_out, text_out = ocr_engine_hocr(ocr_image_out, page_context)
ocr_out = render_hocr_page(hocr_out, page_context)
+2 -8
View File
@@ -6,19 +6,17 @@
from __future__ import annotations
import locale
import logging
import os
import sys
from collections.abc import Sequence
from pathlib import Path
from shutil import copyfileobj
from typing import Union
import pikepdf
from pluggy import PluginManager
from ocrmypdf._defaults import DEFAULT_LANGUAGE, DEFAULT_ROTATE_PAGES_THRESHOLD
from ocrmypdf._defaults import DEFAULT_ROTATE_PAGES_THRESHOLD
from ocrmypdf._exec import unpaper
from ocrmypdf._options import OCROptions
from ocrmypdf.exceptions import (
@@ -74,7 +72,6 @@ def check_options_languages(
raise MissingDependencyError(msg)
def check_options_sidecar(options: OCROptions) -> None:
if options.sidecar == '\0':
if options.output_file == '-':
@@ -117,16 +114,13 @@ def check_options_preprocessing(options: OCROptions) -> None:
raise BadArgsError("--unpaper-args: " + str(e)) from e
def _check_plugin_invariant_options(options: OCROptions) -> None:
check_platform()
check_options_sidecar(options)
check_options_preprocessing(options)
def _check_plugin_options(
options: OCROptions, plugin_manager: PluginManager
) -> None:
def _check_plugin_options(options: OCROptions, plugin_manager: PluginManager) -> None:
plugin_manager.hook.check_options(options=options)
ocr_engine_languages = plugin_manager.hook.get_ocr_engine().languages(options)
check_options_languages(options, ocr_engine_languages)
+36 -25
View File
@@ -43,7 +43,6 @@ import logging
import os
import sys
import threading
from argparse import Namespace
from collections.abc import Iterable, Sequence
from enum import IntEnum
from io import IOBase
@@ -172,8 +171,6 @@ def configure_logging(
return log
def create_options(
*, input_file: PathOrIO, output_file: PathOrIO, parser: ArgumentParser, **kwargs
) -> OCROptions:
@@ -193,30 +190,32 @@ def create_options(
"""
# Prepare kwargs for direct OCROptions construction
options_kwargs = kwargs.copy()
# Set input and output files
options_kwargs['input_file'] = input_file
options_kwargs['output_file'] = output_file
# Handle special stream cases for sidecar
if 'sidecar' in options_kwargs and isinstance(options_kwargs['sidecar'], BinaryIO | IOBase):
if 'sidecar' in options_kwargs and isinstance(
options_kwargs['sidecar'], BinaryIO | IOBase
):
# Keep the stream object as-is - OCROptions can handle it
pass
# Remove None values to let OCROptions use its defaults
options_kwargs = {k: v for k, v in options_kwargs.items() if v is not None}
# Remove any kwargs that aren't OCROptions fields and store in extra_attrs
extra_attrs = {}
ocr_fields = set(OCROptions.model_fields.keys())
# Known extra attributes that should be preserved
known_extra = {'progress_bar', 'plugins'}
for key in list(options_kwargs.keys()):
if key not in ocr_fields and key not in known_extra:
extra_attrs[key] = options_kwargs.pop(key)
# Create OCROptions directly
try:
options = OCROptions(**options_kwargs)
@@ -447,25 +446,29 @@ def _pdf_to_hocr( # noqa: D417
"""
# Prepare kwargs for direct OCROptions construction
options_kwargs = kwargs.copy()
# Set input file and handle special output_folder case
options_kwargs['input_file'] = input_pdf
options_kwargs['output_file'] = '/dev/null' # Placeholder for hOCR pipeline
# Add all the function parameters
for param_name, param_value in locals().items():
if param_name not in {'input_pdf', 'output_folder', 'kwargs', 'plugin_manager', 'plugins'} and param_value is not None:
if (
param_name
not in {'input_pdf', 'output_folder', 'kwargs', 'plugin_manager', 'plugins'}
and param_value is not None
):
options_kwargs[param_name] = param_value
# Handle plugins separately
if plugins:
options_kwargs['plugins'] = plugins
# Remove any kwargs that aren't OCROptions fields and store in extra_attrs
extra_attrs = {'output_folder': output_folder}
ocr_fields = set(OCROptions.model_fields.keys())
known_extra = {'progress_bar', 'plugins'}
for key in list(options_kwargs.keys()):
if key not in ocr_fields and key not in known_extra:
extra_attrs[key] = options_kwargs.pop(key)
@@ -484,7 +487,9 @@ def _pdf_to_hocr( # noqa: D417
if extra_attrs:
options.extra_attrs.update(extra_attrs)
except Exception as e:
raise TypeError(f"Failed to create OCROptions for hOCR pipeline: {e}") from e
raise TypeError(
f"Failed to create OCROptions for hOCR pipeline: {e}"
) from e
return run_hocr_pipeline(options=options, plugin_manager=plugin_manager)
@@ -526,25 +531,29 @@ def _hocr_to_ocr_pdf( # noqa: D417
"""
# Prepare kwargs for direct OCROptions construction
options_kwargs = kwargs.copy()
# Set output file and handle special work_folder case
options_kwargs['input_file'] = '/dev/null' # Placeholder for hOCR to PDF pipeline
options_kwargs['output_file'] = output_file
# Add all the function parameters
for param_name, param_value in locals().items():
if param_name not in {'work_folder', 'output_file', 'kwargs', 'plugin_manager', 'plugins'} and param_value is not None:
if (
param_name
not in {'work_folder', 'output_file', 'kwargs', 'plugin_manager', 'plugins'}
and param_value is not None
):
options_kwargs[param_name] = param_value
# Handle plugins separately
if plugins:
options_kwargs['plugins'] = plugins
# Remove any kwargs that aren't OCROptions fields and store in extra_attrs
extra_attrs = {'work_folder': work_folder}
ocr_fields = set(OCROptions.model_fields.keys())
known_extra = {'progress_bar', 'plugins'}
for key in list(options_kwargs.keys()):
if key not in ocr_fields and key not in known_extra:
extra_attrs[key] = options_kwargs.pop(key)
@@ -563,7 +572,9 @@ def _hocr_to_ocr_pdf( # noqa: D417
if extra_attrs:
options.extra_attrs.update(extra_attrs)
except Exception as e:
raise TypeError(f"Failed to create OCROptions for hOCR to PDF pipeline: {e}") from e
raise TypeError(
f"Failed to create OCROptions for hOCR to PDF pipeline: {e}"
) from e
return run_hocr_to_ocr_pdf_pipeline(
options=options, plugin_manager=plugin_manager
+1 -1
View File
@@ -130,7 +130,7 @@ def generate_pdfa(
output_type = context.options.output_type
if output_type == 'pdfa':
output_type = 'pdfa-2'
ghostscript.generate_pdfa(
pdf_pages=[pdfmark, *pdf_pages],
output_file=output_file,
+2 -2
View File
@@ -263,13 +263,13 @@ class HocrTransform:
def _get_text_direction(self, par):
"""Get the text direction of the paragraph.
Arabic, Hebrew, Persian, are right-to-left languages.
When the paragraph element is None, defaults to left-to-right.
"""
if par is None:
return TextDirection.LTR
return (
TextDirection.RTL
if par.attrib.get('dir', 'ltr') == 'rtl'
+1 -1
View File
@@ -383,7 +383,7 @@ def _get_effective_jbig2_page_group_size(options) -> int:
def extract_images_jbig2(pdf: Pdf, root: Path, options) -> dict[int, list[XrefExt]]:
"""Extract any bitonal image that we think we can improve as JBIG2."""
jbig2_page_group_size = _get_effective_jbig2_page_group_size(options)
jbig2_groups = defaultdict(list)
for pageno, xref_ext in extract_images(pdf, root, options, extract_image_jbig2):
group = pageno // jbig2_page_group_size
-1
View File
@@ -44,7 +44,6 @@ from ocrmypdf.helpers import Resolution, available_cpu_count, pikepdf_enable_mma
from ocrmypdf.pdfinfo.layout import (
LTStateAwareChar,
PdfMinerState,
get_page_analysis,
get_text_boxes,
)
-1
View File
@@ -3,7 +3,6 @@
from __future__ import annotations
import pytest
from pikepdf import Array, Dictionary, Name, NameTree, Pdf
from ocrmypdf._annots import remove_broken_goto_annotations
+16 -18
View File
@@ -1,10 +1,8 @@
"""Test JSON serialization of OCROptions for multiprocessing compatibility."""
import multiprocessing
from pathlib import Path
from io import BytesIO
import pytest
from pathlib import Path
from ocrmypdf._options import OCROptions
@@ -13,7 +11,7 @@ def worker_function(options_json: str) -> str:
"""Worker function that deserializes OCROptions from JSON and returns a result."""
# Reconstruct OCROptions from JSON in worker process
options = OCROptions.model_validate_json_safe(options_json)
# Verify we can access various option types
result = {
'input_file': str(options.input_file),
@@ -24,7 +22,7 @@ def worker_function(options_json: str) -> str:
'fast_web_view': options.fast_web_view,
'extra_attrs_count': len(options.extra_attrs),
}
# Return as JSON string
import json
return json.dumps(result)
@@ -43,14 +41,14 @@ def test_json_serialization_multiprocessing():
deskew=True,
clean=False,
)
# Add some extra attributes
options.extra_attrs['custom_field'] = 'test_value'
options.extra_attrs['numeric_field'] = 42
# Serialize to JSON
options_json = options.model_dump_json_safe()
# Test that we can deserialize in the main process
reconstructed = OCROptions.model_validate_json_safe(options_json)
assert reconstructed.input_file == options.input_file
@@ -62,12 +60,12 @@ def test_json_serialization_multiprocessing():
assert reconstructed.deskew == options.deskew
assert reconstructed.clean == options.clean
assert reconstructed.extra_attrs == options.extra_attrs
# Test multiprocessing with JSON serialization
with multiprocessing.Pool(processes=2) as pool:
# Send the JSON string to worker processes
results = pool.map(worker_function, [options_json, options_json])
# Verify results from worker processes
import json
for result_json in results:
@@ -85,20 +83,20 @@ def test_json_serialization_with_streams():
"""Test JSON serialization with stream objects."""
input_stream = BytesIO(b'fake pdf data')
output_stream = BytesIO()
options = OCROptions(
input_file=input_stream,
output_file=output_stream,
languages=['eng'],
optimize=1,
)
# Serialize to JSON (streams should be converted to placeholders)
options_json = options.model_dump_json_safe()
# Deserialize (streams will be placeholder strings)
reconstructed = OCROptions.model_validate_json_safe(options_json)
# Streams should be converted to placeholder strings
assert reconstructed.input_file == 'stream'
assert reconstructed.output_file == 'stream'
@@ -114,19 +112,19 @@ def test_json_serialization_with_none_values():
languages=['eng'],
# Many fields will be None by default
)
# Serialize to JSON
options_json = options.model_dump_json_safe()
# Deserialize
reconstructed = OCROptions.model_validate_json_safe(options_json)
# Verify None values are preserved (check actual defaults from model)
assert reconstructed.tesseract_timeout == 0.0 # Default value, not None
assert reconstructed.fast_web_view == 1.0 # Default value, not None
assert reconstructed.color_conversion_strategy == "LeaveColorUnchanged" # Default value
assert reconstructed.pdfa_image_compression is None # This one is actually None
# Verify non-None values are preserved
assert reconstructed.input_file == options.input_file
assert reconstructed.output_file == options.output_file
-2
View File
@@ -10,10 +10,8 @@ from shutil import copyfile
import pikepdf
import pytest
from packaging.version import Version
from pikepdf.models.metadata import decode_pdf_date
from ocrmypdf._exec import ghostscript
from ocrmypdf._jobcontext import PdfContext
from ocrmypdf._metadata import metadata_fixup
from ocrmypdf._pipeline import convert_to_pdfa
+1 -1
View File
@@ -13,7 +13,7 @@ import pytest
from ocrmypdf import pdfinfo
from ocrmypdf._exec import tesseract
from ocrmypdf.exceptions import BadArgsError, ExitCode, MissingDependencyError
from ocrmypdf.exceptions import BadArgsError, MissingDependencyError
from .conftest import check_ocrmypdf, run_ocrmypdf_api
+1 -1
View File
@@ -97,7 +97,7 @@ def test_pillow_options():
# Test that max_image_mpixels=0 is valid (validation now in OCROptions)
opts = make_ocr_opts(max_image_mpixels=0)
assert opts.max_image_mpixels == 0
# Test that negative values are rejected
with pytest.raises(ValueError, match="max_image_mpixels must be non-negative"):
make_ocr_opts(max_image_mpixels=-1)