Refactor MDIO metadata handling and test structure.

Author: Altay Sansal

src/mdio/api/accessor.py

@@ -190,28 +190,29 @@ def _validate_store(self, storage_options):
             url=self.url,
             disk_cache=self._disk_cache,
         )
-        self.store = zarr.open(
-            self.url, mode=self.mode, storage_options=storage_options
-        ).store
 
-    def _connect(self):
-        """Open the zarr root."""
         try:
-            if self.mode in {"r", "r+"}:
-                self.root = zarr.open_consolidated(store=self.store, mode=self.mode)
-            elif self.mode == "w":
-                self.root = zarr.open(store=self.store, mode="r+")
-            else:
-                msg = f"Invalid mode: {self.mode}"
-                raise ValueError(msg)
-        except KeyError as e:
+            self.store = zarr.open(
+                self.url, mode=self.mode, storage_options=storage_options
+            ).store
+        except FileNotFoundError as e:
             msg = (
-                f"MDIO file not found or corrupt at {self.store.path}. "
+                f"MDIO file not found or corrupt at {self.url}. "
                 "Please check the URL or ensure it is not a deprecated "
                 "version of MDIO file."
             )
             raise MDIONotFoundError(msg) from e
 
+    def _connect(self):
+        """Open the zarr root."""
+        if self.mode in {"r", "r+"}:
+            self.root = zarr.open_consolidated(store=self.store, mode=self.mode)
+        elif self.mode == "w":
+            self.root = zarr.open(store=self.store, mode="r+")
+        else:
+            msg = f"Invalid mode: {self.mode}"
+            raise ValueError(msg)
+
     def _deserialize_grid(self):
         """Deserialize grid from Zarr metadata."""
         self.grid = Grid.from_zarr(self.root)

src/mdio/api/io_utils.py

@@ -4,7 +4,6 @@
 
 import dask.array as da
 import zarr
-from zarr.storage._utils import normalize_path
 
 
 def process_url(
@@ -69,7 +68,7 @@ def process_url(
     if disk_cache is True:
         url = "::".join(["simplecache", url])
 
-    return normalize_path(url)
+    return url
 
 
 def open_zarr_array(group_handle: zarr.Group, name: str) -> zarr.Array:

src/mdio/segy/helpers_segy.py

@@ -29,7 +29,7 @@ def create_zarr_hierarchy(root_group: "Group", overwrite: bool) -> "Group":
         root_group.create_group(name="metadata", overwrite=overwrite)
     except ContainsGroupError as e:
         msg = (
-            f"An MDIO file with data already exists at {root_group.store.path}. "
+            f"An MDIO file with data already exists at {root_group.store}. "
             "If this is intentional, please specify 'overwrite=True'."
         )
         raise MDIOAlreadyExistsError(msg) from e

tests/integration/test_segy_import_export.py

@@ -259,7 +259,7 @@ def test_3d_import(segy_input, zarr_tmp, index_bytes, index_names):
     """Test importing a SEG-Y file to MDIO."""
     segy_to_mdio(
         segy_path=segy_input.__str__(),
-        mdio_path_or_buffer=zarr_tmp.__str__(),
+        mdio_path_or_buffer=(zarr_tmp / "teapot.mdio").__str__(),
         index_bytes=index_bytes,
         index_names=index_names,
         overwrite=True,
@@ -272,13 +272,13 @@ class TestReader:
 
     def test_meta_read(self, zarr_tmp):
         """Metadata reading tests."""
-        mdio = MDIOReader(zarr_tmp.__str__())
+        mdio = MDIOReader((zarr_tmp / "teapot.mdio").__str__())
         assert mdio.binary_header["samples_per_trace"] == 1501
         assert mdio.binary_header["sample_interval"] == 2000
 
     def test_grid(self, zarr_tmp):
         """Grid reading tests."""
-        mdio = MDIOReader(zarr_tmp.__str__())
+        mdio = MDIOReader((zarr_tmp / "teapot.mdio").__str__())
         grid = mdio.grid
 
         assert grid.select_dim("inline") == Dimension(range(1, 346), "inline")
@@ -287,7 +287,7 @@ def test_grid(self, zarr_tmp):
 
     def test_get_data(self, zarr_tmp):
         """Data retrieval tests."""
-        mdio = MDIOReader(zarr_tmp.__str__())
+        mdio = MDIOReader((zarr_tmp / "teapot.mdio").__str__())
 
         assert mdio.shape == (345, 188, 1501)
         assert mdio[0, :, :].shape == (188, 1501)
@@ -296,7 +296,7 @@ def test_get_data(self, zarr_tmp):
 
     def test_inline(self, zarr_tmp):
         """Read and compare every 75 inlines' mean and std. dev."""
-        mdio = MDIOReader(zarr_tmp.__str__())
+        mdio = MDIOReader((zarr_tmp / "teapot.mdio").__str__())
 
         inlines = mdio[::75, :, :]
         mean, std = inlines.mean(), inlines.std()
@@ -305,7 +305,7 @@ def test_inline(self, zarr_tmp):
 
     def test_crossline(self, zarr_tmp):
         """Read and compare every 75 crosslines' mean and std. dev."""
-        mdio = MDIOReader(zarr_tmp.__str__())
+        mdio = MDIOReader((zarr_tmp / "teapot.mdio").__str__())
 
         xlines = mdio[:, ::75, :]
         mean, std = xlines.mean(), xlines.std()
@@ -314,7 +314,7 @@ def test_crossline(self, zarr_tmp):
 
     def test_zslice(self, zarr_tmp):
         """Read and compare every 225 z-slices' mean and std. dev."""
-        mdio = MDIOReader(zarr_tmp.__str__())
+        mdio = MDIOReader((zarr_tmp / "teapot.mdio").__str__())
 
         slices = mdio[:, :, ::225]
         mean, std = slices.mean(), slices.std()
@@ -329,7 +329,7 @@ class TestExport:
     def test_3d_export(self, zarr_tmp, segy_export_tmp):
         """Test 3D export to IBM and IEEE."""
         mdio_to_segy(
-            mdio_path_or_buffer=zarr_tmp.__str__(),
+            mdio_path_or_buffer=(zarr_tmp / "teapot.mdio").__str__(),
             output_segy_path=segy_export_tmp.__str__(),
         )
 

tests/unit/conftest.py

@@ -4,12 +4,12 @@
 
 from datetime import datetime
 from importlib import metadata
+from pathlib import Path
 
 import numpy as np
 import pytest
 import zarr
 from numpy.typing import NDArray
-from zarr import Group
 from zarr import consolidate_metadata
 from zarr import open_group
 
@@ -30,11 +30,9 @@
 
 
 @pytest.fixture(scope="module")
-def mock_root_group(tmp_path_factory) -> Group:
+def mock_mdio_path(tmp_path_factory):
     """Make a mocked MDIO store for writing."""
-    zarr.config.set({"default_zarr_format": 2, "write_empty_chunks": False})
-    tmp_dir = tmp_path_factory.mktemp("mdio")
-    return open_group(tmp_dir.name, mode="w")
+    return tmp_path_factory.mktemp("mdio") / "mock.mdio"
 
 
 @pytest.fixture
@@ -75,16 +73,19 @@ def mock_data(mock_coords):
 
 @pytest.fixture
 def mock_mdio(
-    mock_root_group: Group,
+    mock_mdio_path: Path,
     mock_dimensions: list[Dimension],
     mock_coords: tuple[NDArray],
     mock_data: NDArray,
     mock_text: list[str],
     mock_bin: dict[str, int],
 ):
     """This mocks most of mdio.converters.segy in memory."""
-    zarr_root = create_zarr_hierarchy(
-        root_group=mock_root_group,
+    zarr.config.set({"default_zarr_format": 2, "write_empty_chunks": False})
+    zarr_root = open_group(mock_mdio_path, mode="w")
+
+    create_zarr_hierarchy(
+        root_group=zarr_root,
         overwrite=True,
     )
 
@@ -114,6 +115,7 @@ def mock_mdio(
         shape=grid.shape[:-1],
         chunks=grid.shape[:-1],
         chunk_key_encoding={"name": "v2", "separator": "/"},
+        overwrite=True,
     )
     live_mask_arr[...] = grid.live_mask[...]
 
@@ -135,34 +137,39 @@ def mock_mdio(
 
     data_arr = data_grp.create_array(
         "chunked_012",
-        data=mock_data,
+        dtype=mock_data.dtype,
+        shape=mock_data.shape,
         chunk_key_encoding={"name": "v2", "separator": "/"},
+        overwrite=True,
     )
+    data_arr[...] = mock_data
 
-    metadata_grp.create_array(
-        data=il_grid * xl_grid,
+    metadata_arr = metadata_grp.create_array(
         name="_".join(["chunked_012", "trace_headers"]),
+        dtype=il_grid.dtype,
         shape=grid.shape[:-1],  # Same spatial shape as data
         chunks=data_arr.chunks[:-1],  # Same spatial chunks as data
         chunk_key_encoding={"name": "v2", "separator": "/"},
+        overwrite=True,
     )
+    metadata_arr[...] = il_grid * xl_grid
 
-    consolidate_metadata(mock_root_group.store)
+    consolidate_metadata(zarr_root.store)
 
-    return zarr_root
+    return mock_mdio_path
 
 
 @pytest.fixture
-def mock_reader(mock_mdio: Group) -> MDIOReader:
+def mock_reader(mock_mdio: Path) -> MDIOReader:
     """Reader that points to the mocked data to be used later."""
-    return MDIOReader(mock_mdio.store.path)
+    return MDIOReader(mock_mdio.__str__())
 
 
 @pytest.fixture
-def mock_reader_cached(mock_mdio: Group) -> MDIOReader:
+def mock_reader_cached(mock_mdio: Path) -> MDIOReader:
     """Reader that points to the mocked data to be used later. (with local caching)."""
     return MDIOReader(
-        mock_mdio.store.path,
+        mock_mdio.__str__(),
         disk_cache=True,
         storage_options={"simplecache": {"cache_storage": "./mdio_test_cache"}},
     )

tests/unit/test_accessor.py

@@ -132,6 +132,6 @@ def test_wrong_index(self, mock_reader: MDIOReader) -> None:
 
     def test_mdio_exists(self, mock_reader: MDIOReader) -> None:
         """MDIO doesn't exist or corrupt."""
-        mock_store = mock_reader.store
+        mock_root = mock_reader.root
         with pytest.raises(MDIOAlreadyExistsError):
-            create_zarr_hierarchy(mock_store, overwrite=False)
+            create_zarr_hierarchy(mock_root, overwrite=False)

tests/unit/test_compat.py

@@ -27,18 +27,20 @@
 def update_mdio_for_version_0_7_4(root):
     """Update MDIO metadata to mimic version 0.7.4."""
     # Update binary header revision keys
-    bin_hdr = root.metadata.attrs[BINARY_HEADER_KEY]
+    bin_hdr = root["metadata"].attrs[BINARY_HEADER_KEY]
     bin_hdr["SEGYRevision"] = bin_hdr.pop("segy_revision_major")
     bin_hdr["SEGYRevisionMinor"] = bin_hdr.pop("segy_revision_minor")
-    root.metadata.attrs[BINARY_HEADER_KEY] = bin_hdr
+    root["metadata"].attrs[BINARY_HEADER_KEY] = bin_hdr
 
     # Remove trace headers past field 232 (pre-0.8 schema)
-    orig_hdr = root.metadata[CHUNKED_TRACE_HEADERS_KEY]
+    orig_hdr = root["metadata"][CHUNKED_TRACE_HEADERS_KEY]
     new_dtype = np.dtype(orig_hdr.dtype.descr[:-1])
-    new_hdr = zarr.zeros_like(orig_hdr, dtype=new_dtype)
-    root.metadata.create_dataset(
-        CHUNKED_TRACE_HEADERS_KEY,
-        data=new_hdr,
+    root["metadata"].create_array(
+        name=CHUNKED_TRACE_HEADERS_KEY,
+        shape=orig_hdr.shape,
+        dtype=new_dtype,
+        chunks=orig_hdr.chunks,
+        chunk_key_encoding={"name": "v2", "separator": "/"},
         overwrite=True,
     )
     zarr.consolidate_metadata(root.store)