From dbd3c937578a238bb3eb70589100d41839e0dc2d Mon Sep 17 00:00:00 2001
From: "James R. Barlow" <james@purplerock.ca>
Date: Sun, 10 Nov 2024 02:05:57 -0800
Subject: [PATCH] Fix issue with unpickling HOCRResult

Fixes [Bug]: HOCRResult.from_json() not unpickling correctly #1427
---
 src/ocrmypdf/_pipelines/_common.py | 45 +++++++++++++-----------------
 tests/test_api.py                  | 30 +++++++++++++++++++-
 2 files changed, 48 insertions(+), 27 deletions(-)

diff --git a/src/ocrmypdf/_pipelines/_common.py b/src/ocrmypdf/_pipelines/_common.py
index cc134c52..23d1df8f 100644
--- a/src/ocrmypdf/_pipelines/_common.py
+++ b/src/ocrmypdf/_pipelines/_common.py
@@ -104,6 +104,23 @@ class PageResult(NamedTuple):
     """Orientation correction in degrees."""
 
 
+class HOCRResultEncoder(json.JSONEncoder):
+    def default(self, obj):
+        if isinstance(obj, Path):
+            return {'Path': str(obj)}
+        return super().default(obj)
+
+
+class HOCRResultDecoder(json.JSONDecoder):
+    def __init__(self, *args, **kwargs):
+        super().__init__(object_hook=self.dict_to_object, *args, **kwargs)
+
+    def dict_to_object(self, d):
+        if 'Path' in d:
+            return Path(d['Path'])
+        return d
+
+
 @dataclass
 class HOCRResult:
     """Result when hOCR is finished processing."""
@@ -123,38 +140,14 @@ class HOCRResult:
     orientation_correction: int = 0
     """Orientation correction in degrees."""
 
-    def __getstate__(self):
-        """Return state values to be pickled."""
-        return {
-            k: (
-                ('Path://' + str(v))
-                if k in ('pdf_page_from_image', 'hocr', 'textpdf') and v is not None
-                else v
-            )
-            for k, v in self.__dict__.items()
-        }
-
-    def __setstate__(self, state):
-        """Restore state from the unpickled state values."""
-        self.__dict__.update(
-            {
-                k: (
-                    Path(v.removeprefix('Path://'))
-                    if k in ('pdf_page_from_image', 'hocr', 'textpdf') and v is not None
-                    else v
-                )
-                for k, v in state.items()
-            }
-        )
-
     @classmethod
     def from_json(cls, json_str: str) -> HOCRResult:
         """Create an instance from a dict."""
-        return cls(**json.loads(json_str))
+        return cls(**json.loads(json_str, cls=HOCRResultDecoder))
 
     def to_json(self) -> str:
         """Serialize to a JSON string."""
-        return json.dumps(self.__getstate__())
+        return json.dumps(self.__dict__, cls=HOCRResultEncoder)
 
 
 def configure_debug_logging(
diff --git a/tests/test_api.py b/tests/test_api.py
index f5eea148..d6f34fcd 100644
--- a/tests/test_api.py
+++ b/tests/test_api.py
@@ -3,6 +3,7 @@
 
 from __future__ import annotations
 
+import pickle
 from io import BytesIO
 from pathlib import Path
 
@@ -10,6 +11,7 @@ import pytest
 from pdfminer.high_level import extract_text
 
 import ocrmypdf
+import ocrmypdf._pipelines
 import ocrmypdf.api
 
 
@@ -35,7 +37,7 @@ def test_sidecar_stringio(resources: Path, outdir: Path, outpdf: Path):
         resources / 'ccitt.pdf',
         outpdf,
         plugins=['tests/plugins/tesseract_cache.py'],
-        sidecar=s
+        sidecar=s,
     )
     s.seek(0)
     assert b'the' in s.getvalue()
@@ -75,3 +77,29 @@ def test_hocr_to_pdf_api(resources: Path, outdir: Path, outpdf: Path):
     text = extract_text(outpdf)
     assert 'hocr' in text and 'the' not in text
 
+
+def test_hocr_result_json():
+    result = ocrmypdf._pipelines._common.HOCRResult(
+        pageno=1,
+        pdf_page_from_image=Path('a'),
+        hocr=Path('b'),
+        textpdf=Path('c'),
+        orientation_correction=180,
+    )
+    assert (
+        result.to_json()
+        == '{"pageno": 1, "pdf_page_from_image": {"Path": "a"}, "hocr": {"Path": "b"}, '
+        '"textpdf": {"Path": "c"}, "orientation_correction": 180}'
+    )
+    assert ocrmypdf._pipelines._common.HOCRResult.from_json(result.to_json()) == result
+
+
+def test_hocr_result_pickle():
+    result = ocrmypdf._pipelines._common.HOCRResult(
+        pageno=1,
+        pdf_page_from_image=Path('a'),
+        hocr=Path('b'),
+        textpdf=Path('c'),
+        orientation_correction=180,
+    )
+    assert result == pickle.loads(pickle.dumps(result))