feat: Implement JSON serialization for OCROptions with safe handling of Path and stream objects
Co-authored-by: aider (openrouter/anthropic/claude-sonnet-4) <aider@aider.chat>
This commit is contained in:
James R. Barlow
aider
+35
-44
@@ -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
|
||||
|
||||
@@ -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.
|
||||
|
||||
Reference in New Issue
Block a user