Add ObstoreMemCacheReader (#2)

Author: maxrjones

src/obspec_utils/__init__.py

@@ -1,4 +1,4 @@
 from ._version import __version__
-from .file_handlers import ObstoreReader
+from .file_handlers import ObstoreMemCacheReader, ObstoreReader
 
-__all__ = ["__version__", "ObstoreReader"]
+__all__ = ["__version__", "ObstoreMemCacheReader", "ObstoreReader"]

src/obspec_utils/file_handlers.py

@@ -8,6 +8,8 @@
     from obstore import ReadableFile
     from obstore.store import ObjectStore
 
+from obstore.store import MemoryStore
+
 
 class ObstoreReader:
     _reader: ReadableFile
@@ -40,3 +42,30 @@ def seek(self, offset: int, whence: int = 0, /):
 
     def tell(self) -> int:
         return self._reader.tell()
+
+
+class ObstoreMemCacheReader(ObstoreReader):
+    _reader: ReadableFile
+    _memstore: MemoryStore
+
+    def __init__(self, store: ObjectStore, path: str) -> None:
+        """
+        Create an obstore file reader that caches the specified path
+        in a MemoryStore then performs reads from the file in memory.
+
+        This reader loads the entire file into memory first, which can be beneficial
+        for files that will be read multiple times or when you want to avoid repeated
+        network requests to the original store.
+
+        Parameters
+        ----------
+        store
+            [ObjectStore][obstore.store.ObjectStore] for reading the file.
+        path
+            The path to the file within the store. This should not include the prefix.
+        """
+        self._memstore = MemoryStore()
+        buffer = store.get(path).bytes()
+        self._memstore.put(path, buffer)
+
+        self._reader = obs.open_reader(self._memstore, path)

tests/test_xarray.py

@@ -1,5 +1,5 @@
 import xarray as xr
-from obspec_utils import ObstoreReader
+from obspec_utils import ObstoreMemCacheReader, ObstoreReader
 from obstore.store import LocalStore
 
 
@@ -8,3 +8,55 @@ def test_local_reader(local_netcdf4_file) -> None:
     reader = ObstoreReader(store=LocalStore(), path=local_netcdf4_file)
     ds_obstore = xr.open_dataset(reader, engine="h5netcdf")
     xr.testing.assert_allclose(ds_fsspec, ds_obstore)
+
+
+def test_memcache_reader(local_netcdf4_file) -> None:
+    """Test that ObstoreMemCacheReader works with xarray."""
+    ds_fsspec = xr.open_dataset(local_netcdf4_file, engine="h5netcdf")
+    reader = ObstoreMemCacheReader(store=LocalStore(), path=local_netcdf4_file)
+    ds_obstore = xr.open_dataset(reader, engine="h5netcdf")
+    xr.testing.assert_allclose(ds_fsspec, ds_obstore)
+
+
+def test_memcache_reader_interface(local_netcdf4_file) -> None:
+    """Test that ObstoreMemCacheReader implements the same interface as ObstoreReader."""
+    store = LocalStore()
+    regular_reader = ObstoreReader(store=store, path=local_netcdf4_file)
+    memcache_reader = ObstoreMemCacheReader(store=store, path=local_netcdf4_file)
+
+    # Test readall
+    data_regular = regular_reader.readall()
+    data_memcache = memcache_reader.readall()
+    assert data_regular == data_memcache
+    assert isinstance(data_memcache, bytes)
+
+
+def test_memcache_reader_multiple_reads(local_netcdf4_file) -> None:
+    """Test that ObstoreMemCacheReader can perform multiple reads."""
+    store = LocalStore()
+    reader = ObstoreMemCacheReader(store=store, path=local_netcdf4_file)
+
+    # Read the first 100 bytes
+    chunk1 = reader.read(100)
+    assert len(chunk1) == 100
+    assert isinstance(chunk1, bytes)
+
+    # Read the next 100 bytes
+    chunk2 = reader.read(100)
+    assert len(chunk2) == 100
+    assert isinstance(chunk2, bytes)
+
+    # The two chunks should be different (different parts of the file)
+    assert chunk1 != chunk2
+
+    # Test tell
+    position = reader.tell()
+    assert position == 200
+
+    # Test seek
+    reader.seek(0)
+    assert reader.tell() == 0
+
+    # Re-reading from the beginning should give us the same data
+    chunk1_again = reader.read(100)
+    assert chunk1 == chunk1_again