diff --git a/src/ocrmypdf/_jobcontext.py b/src/ocrmypdf/_jobcontext.py index 293fc99c..a6495835 100644 --- a/src/ocrmypdf/_jobcontext.py +++ b/src/ocrmypdf/_jobcontext.py @@ -38,15 +38,10 @@ class PdfContext: # Handle both OCROptions and Namespace during transition if isinstance(options, OCROptions): self.options = options - self._namespace_options = options.to_namespace() else: # Convert Namespace to OCROptions self.options = OCROptions.from_namespace(options) - self._namespace_options = self.options.to_namespace() - # Ensure lossless_reconstruction is available on the namespace - if not hasattr(self._namespace_options, 'lossless_reconstruction'): - self._namespace_options.lossless_reconstruction = self.options.lossless_reconstruction self.work_folder = work_folder self.origin = origin self.pdfinfo = pdfinfo @@ -79,8 +74,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 Namespace options instead of OCROptions for pickle compatibility - in multiprocessing scenarios. + Note: Uses OCROptions with JSON serialization for multiprocessing compatibility. """ origin: Path #: The filename of the original input file. @@ -91,8 +85,8 @@ class PageContext: def __init__(self, pdf_context: PdfContext, pageno): self.work_folder = pdf_context.work_folder self.origin = pdf_context.origin - # Always use Namespace for PageContext to avoid pickling issues - self.options = pdf_context._namespace_options + # Store OCROptions directly instead of Namespace + self.options = pdf_context.options self.pageno = pageno self.pageinfo = pdf_context.pdfinfo[pageno] self.plugin_manager = pdf_context.plugin_manager @@ -110,43 +104,40 @@ class PageContext: def __getstate__(self): state = self.__dict__.copy() - # Ensure we only pickle the Namespace, not any Pydantic objects - # Create a completely new Namespace to avoid any contamination - from argparse import Namespace - import os + # Use JSON serialization instead of Namespace + try: + options_json = self.options.model_dump_json_safe() + state['options_json'] = options_json + # Remove the OCROptions object to avoid pickle issues + del state['options'] + except Exception: + # 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).items(): - if key.startswith('_'): - continue - try: - import pickle - - pickle.dumps(value) - setattr(clean_options, key, value) - except TypeError: - continue - # Set lossless_reconstruction if it exists, otherwise compute it - if hasattr(self.options, 'lossless_reconstruction'): - clean_options.lossless_reconstruction = self.options.lossless_reconstruction - else: - # Compute lossless_reconstruction for Namespace objects - clean_options.lossless_reconstruction = not any([ - getattr(self.options, 'deskew', False), - getattr(self.options, 'clean_final', False), - getattr(self.options, 'force_ocr', False), - getattr(self.options, 'remove_background', False), - ]) - state['options'] = clean_options - - # Handle stream inputs - if hasattr(state['options'], 'input_file'): - if not isinstance(state['options'].input_file, str | bytes | os.PathLike): - state['options'].input_file = 'stream' - if hasattr(state['options'], 'output_file'): - if not isinstance(state['options'].output_file, str | bytes | os.PathLike): - state['options'].output_file = 'stream' + clean_options = Namespace() + for key, value in vars(self.options.to_namespace()).items(): + if key.startswith('_'): + continue + try: + import pickle + pickle.dumps(value) + setattr(clean_options, key, value) + except TypeError: + continue + state['options'] = clean_options # Remove any potential references to Pydantic objects state.pop('_pdf_context', None) return state + + 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 diff --git a/src/ocrmypdf/_options.py b/src/ocrmypdf/_options.py index 22274926..be6f169b 100644 --- a/src/ocrmypdf/_options.py +++ b/src/ocrmypdf/_options.py @@ -5,6 +5,7 @@ from __future__ import annotations +import json import logging import os import unicodedata @@ -377,6 +378,74 @@ class OCROptions(BaseModel): self.extra_attrs['lossless_reconstruction'] = lossless return self + def model_dump_json_safe(self) -> str: + """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 hasattr(value, 'read') or hasattr(value, 'write'): + # Stream object - replace with placeholder + return {'__type__': 'Stream', 'value': 'stream'} + elif isinstance(value, (list, tuple)): + return [_serialize_value(item) for item in value] + elif isinstance(value, dict): + 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: + if value['__type__'] == 'Path': + return Path(value['value']) + elif value['__type__'] == 'Stream': + # For streams, we'll use a placeholder string + return value['value'] + else: + return value['value'] + elif isinstance(value, list): + return [_deserialize_value(item) for item in value] + elif isinstance(value, dict): + 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( extra="forbid", # Force use of extra_attrs for unknown fields arbitrary_types_allowed=True, # Allow BinaryIO, Path, etc.