Fix: Recursion when pickling (#37)

* Fix: Pickling readers and store wrappers

* Test __get_attr__

* Test composition

Author: maxrjones

src/obspec_utils/cache.py

@@ -109,8 +109,18 @@ def __getattr__(self, name: str) -> Any:
         """Forward unknown attributes to the underlying store.
 
         This ensures CachingReadableStore is transparent for any additional
-        methods the underlying store may have (e.g., prefix, url attributes).
+        public methods or attributes the underlying store may have.
+
+        Note: Private attributes (starting with '_') are not forwarded.
         """
+        if name.startswith("_"):
+            raise AttributeError(
+                f"'{type(self).__name__}' object has no attribute '{name}'"
+            )
+        if "_store" not in self.__dict__:
+            raise AttributeError(
+                f"'{type(self).__name__}' object has no attribute '{name}'"
+            )
         return getattr(self._store, name)
 
     def __reduce__(self):

src/obspec_utils/obspec.py

@@ -429,6 +429,9 @@ def __init__(
             would require more requests than this, request sizes are increased to
             fit within this limit.
         """
+        self._store = store
+        self._path = path
+
         # Determine file size if not provided
         if file_size is None:
             if hasattr(store, "head") and callable(store.head):

src/obspec_utils/splitting.py

@@ -131,7 +131,21 @@ def __init__(
         self._max_concurrent_requests = max_concurrent_requests
 
     def __getattr__(self, name: str) -> Any:
-        """Forward unknown attributes to the underlying store."""
+        """Forward unknown attributes to the underlying store.
+
+        This ensures SplittingReadableStore is transparent for any additional
+        public methods or attributes the underlying store may have.
+
+        Note: Private attributes (starting with '_') are not forwarded.
+        """
+        if name.startswith("_"):
+            raise AttributeError(
+                f"'{type(self).__name__}' object has no attribute '{name}'"
+            )
+        if "_store" not in self.__dict__:
+            raise AttributeError(
+                f"'{type(self).__name__}' object has no attribute '{name}'"
+            )
         return getattr(self._store, name)
 
     def _get_file_size(self, path: str) -> int | None:

tests/mocks.py

@@ -136,3 +136,65 @@ async def get_ranges_async(self, path, *, starts, ends=None, lengths=None):
         if ends is None:
             ends = [s + ln for s, ln in zip(starts, lengths)]
         return [self._data[s:e] for s, e in zip(starts, ends)]
+
+
+class PicklableStore:
+    """A picklable store supporting multiple paths, for testing pickle support.
+
+    Unlike obstore's MemoryStore (Rust-backed, not picklable), this pure-Python
+    store can be pickled and unpickled, allowing tests to verify pickle support
+    for wrappers like CachingReadableStore and reader classes.
+    """
+
+    def __init__(self, data: dict[str, bytes] | None = None):
+        self._data = data if data is not None else {}
+
+    def put(self, path: str, data: bytes) -> None:
+        self._data[path] = data
+
+    def head(self, path: str) -> dict:
+        return {"size": len(self._data[path])}
+
+    def get(self, path: str, *, options=None):
+        return MockGetResult(self._data[path])
+
+    async def get_async(self, path: str, *, options=None):
+        return MockGetResultAsync(self._data[path])
+
+    def get_range(
+        self,
+        path: str,
+        *,
+        start: int,
+        end: int | None = None,
+        length: int | None = None,
+    ):
+        data = self._data[path]
+        if length is not None:
+            return data[start : start + length]
+        elif end is not None:
+            return data[start:end]
+        return data[start:]
+
+    async def get_range_async(
+        self,
+        path: str,
+        *,
+        start: int,
+        end: int | None = None,
+        length: int | None = None,
+    ):
+        return self.get_range(path, start=start, end=end, length=length)
+
+    def get_ranges(self, path: str, *, starts, ends=None, lengths=None):
+        if lengths is not None:
+            return [
+                self._data[path][start : start + length]
+                for start, length in zip(starts, lengths)
+            ]
+        elif ends is not None:
+            return [self._data[path][s:e] for s, e in zip(starts, ends)]
+        raise ValueError("Must provide ends or lengths")
+
+    async def get_ranges_async(self, path: str, *, starts, ends=None, lengths=None):
+        return self.get_ranges(path, starts=starts, ends=ends, lengths=lengths)

tests/test_cache.py

@@ -10,6 +10,8 @@
 from obspec_utils.cache import CachingReadableStore
 from obspec_utils.registry import ObjectStoreRegistry
 
+from .mocks import PicklableStore
+
 
 class TestCachingReadableStore:
     """Tests for basic caching functionality."""
@@ -392,84 +394,6 @@ def test_forwards_unknown_attributes(self):
         assert hasattr(cached, "delete")  # MemoryStore has delete
 
 
-class PicklableStore:
-    """A simple picklable store for testing pickle support.
-
-    MemoryStore from obstore is Rust-backed and not picklable.
-    This pure-Python store allows testing CachingReadableStore's pickle support.
-    """
-
-    def __init__(self, data: dict[str, bytes] | None = None):
-        self._data = data or {}
-
-    def put(self, path: str, data: bytes) -> None:
-        self._data[path] = data
-
-    def get(self, path: str, *, options=None):
-        return _PicklableGetResult(self._data[path])
-
-    async def get_async(self, path: str, *, options=None):
-        return _PicklableGetResultAsync(self._data[path])
-
-    def get_range(
-        self,
-        path: str,
-        *,
-        start: int,
-        end: int | None = None,
-        length: int | None = None,
-    ):
-        data = self._data[path]
-        if end is not None:
-            return data[start:end]
-        elif length is not None:
-            return data[start : start + length]
-        return data[start:]
-
-    async def get_range_async(
-        self,
-        path: str,
-        *,
-        start: int,
-        end: int | None = None,
-        length: int | None = None,
-    ):
-        return self.get_range(path, start=start, end=end, length=length)
-
-    def get_ranges(self, path: str, *, starts, ends=None, lengths=None):
-        if ends is not None:
-            return [self._data[path][s:e] for s, e in zip(starts, ends)]
-        elif lengths is not None:
-            return [
-                self._data[path][start : start + length]
-                for start, length in zip(starts, lengths)
-            ]
-        raise ValueError("Must provide ends or lengths")
-
-    async def get_ranges_async(self, path: str, *, starts, ends=None, lengths=None):
-        return self.get_ranges(path, starts=starts, ends=ends, lengths=lengths)
-
-
-class _PicklableGetResult:
-    """Mock GetResult for PicklableStore."""
-
-    def __init__(self, data: bytes):
-        self._data = data
-
-    def buffer(self):
-        return self._data
-
-
-class _PicklableGetResultAsync:
-    """Mock async GetResult for PicklableStore."""
-
-    def __init__(self, data: bytes):
-        self._data = data
-
-    async def buffer_async(self):
-        return self._data
-
-
 class TestPickling:
     """Tests for pickling support (needed for multiprocessing/distributed)."""
 

tests/test_readers.py

@@ -3,6 +3,7 @@
 EagerStoreReader, and ParallelStoreReader.
 """
 
+import pickle
 from io import BytesIO
 
 import pytest
@@ -14,6 +15,8 @@
     ParallelStoreReader,
 )
 
+from .mocks import PicklableStore
+
 
 ALL_READERS = [BufferedStoreReader, EagerStoreReader, ParallelStoreReader]
 
@@ -386,3 +389,86 @@ def test_reader_seek_invalid_whence_raises(ReaderClass):
 
     with pytest.raises(ValueError):
         reader.seek(0, 3)
+
+
+# =============================================================================
+# Pickling tests
+# =============================================================================
+
+
+@pytest.mark.parametrize("ReaderClass", ALL_READERS)
+def test_reader_pickle_roundtrip(ReaderClass):
+    """Reader can be pickled and unpickled."""
+    store = PicklableStore()
+    store.put("test.txt", b"hello world")
+
+    reader = ReaderClass(store, "test.txt")
+
+    pickled = pickle.dumps(reader)
+    restored = pickle.loads(pickled)
+
+    assert isinstance(restored, ReaderClass)
+
+
+@pytest.mark.parametrize("ReaderClass", ALL_READERS)
+def test_reader_pickle_preserves_path(ReaderClass):
+    """Unpickled reader preserves the file path."""
+    store = PicklableStore()
+    store.put("test.txt", b"hello world")
+
+    reader = ReaderClass(store, "test.txt")
+    restored = pickle.loads(pickle.dumps(reader))
+
+    assert restored._path == "test.txt"
+
+
+@pytest.mark.parametrize("ReaderClass", ALL_READERS)
+def test_reader_pickle_restored_is_functional(ReaderClass):
+    """Restored reader can read data."""
+    store = PicklableStore()
+    store.put("test.txt", b"hello world")
+
+    reader = ReaderClass(store, "test.txt")
+    restored = pickle.loads(pickle.dumps(reader))
+
+    # Should be able to read data
+    assert restored.read(5) == b"hello"
+    assert restored.read(6) == b" world"
+
+
+@pytest.mark.parametrize("ReaderClass", ALL_READERS)
+def test_reader_pickle_preserves_position(ReaderClass):
+    """Unpickled reader preserves the current position."""
+    store = PicklableStore()
+    store.put("test.txt", b"hello world")
+
+    reader = ReaderClass(store, "test.txt")
+    reader.read(5)  # Move position to 5
+    assert reader.tell() == 5
+
+    restored = pickle.loads(pickle.dumps(reader))
+
+    assert restored.tell() == 5
+    assert restored.read(6) == b" world"
+
+
+@pytest.mark.parametrize("ReaderClass", ALL_READERS)
+def test_reader_pickle_multiple_protocols(ReaderClass):
+    """Pickling works with different pickle protocols.
+
+    Note: EagerStoreReader uses BytesIO which requires protocol >= 2.
+    """
+    store = PicklableStore()
+    store.put("test.txt", b"hello world")
+
+    reader = ReaderClass(store, "test.txt")
+
+    # BytesIO (used by EagerStoreReader) requires protocol >= 2
+    min_protocol = 2 if ReaderClass == EagerStoreReader else 0
+
+    for protocol in range(min_protocol, pickle.HIGHEST_PROTOCOL + 1):
+        pickled = pickle.dumps(reader, protocol=protocol)
+        restored = pickle.loads(pickled)
+
+        restored.seek(0)
+        assert restored.read() == b"hello world"