Implement Pydantic grid override definitions

Author: BrianMichell

src/mdio/__init__.py

@@ -10,6 +10,7 @@
 from mdio.converters import segy_to_mdio
 from mdio.optimize.access_pattern import OptimizedAccessPatternConfig
 from mdio.optimize.access_pattern import optimize_access_patterns
+from mdio.segy.geometry import GridOverrides
 
 try:
     __version__ = metadata.version("multidimio")
@@ -19,6 +20,7 @@
 
 __all__ = [
     "__version__",
+    "GridOverrides",
     "open_mdio",
     "to_mdio",
     "mdio_to_segy",

src/mdio/converters/segy.py

@@ -37,6 +37,7 @@
 from mdio.ingestion.segy.validation import _validate_spec_in_template
 from mdio.segy import blocked_io
 from mdio.segy.file import get_segy_file_info
+from mdio.segy.geometry import GridOverrides
 from mdio.segy.utilities import get_grid_plan
 
 if TYPE_CHECKING:
@@ -128,7 +129,7 @@ def filtered_add_coordinate(  # noqa: ANN202
 
 def _update_template_from_grid_overrides(
     template: AbstractDatasetTemplate,
-    grid_overrides: dict[str, Any] | None,
+    grid_overrides: GridOverrides | None,
     segy_dimensions: list[Dimension],
     full_chunk_shape: tuple[int, ...],
     chunk_size: tuple[int, ...],
@@ -178,30 +179,29 @@ def _update_template_from_grid_overrides(
 
     # If using NonBinned override, expose non-binned dims as logical coordinates on the template instance
     # and patch _add_coordinates to skip adding them as 1D dimension coordinates
-    if grid_overrides and "NonBinned" in grid_overrides and "non_binned_dims" in grid_overrides:
-        non_binned_dims = tuple(grid_overrides["non_binned_dims"])
-        if non_binned_dims:
-            logger.debug(
-                "NonBinned grid override: exposing non-binned dims as coordinates: %s",
-                non_binned_dims,
-            )
-            # Append any missing names; keep existing order and avoid duplicates
-            existing = set(template.coordinate_names)
-            to_add = tuple(n for n in non_binned_dims if n not in existing)
-            if to_add:
-                template._logical_coord_names = template._logical_coord_names + to_add
-
-            # Patch _add_coordinates to skip adding non-binned dims as 1D dimension coordinates
-            # This prevents them from being added with wrong dimensions (e.g., just "trace")
-            # They will be added later by build_dataset with full spatial_dimension_names
-            _patch_add_coordinates_for_non_binned(template, set(non_binned_dims))
+    if grid_overrides is not None and grid_overrides.non_binned and grid_overrides.non_binned_dims:
+        non_binned_dims = tuple(grid_overrides.non_binned_dims)
+        logger.debug(
+            "NonBinned grid override: exposing non-binned dims as coordinates: %s",
+            non_binned_dims,
+        )
+        # Append any missing names; keep existing order and avoid duplicates
+        existing = set(template.coordinate_names)
+        to_add = tuple(n for n in non_binned_dims if n not in existing)
+        if to_add:
+            template._logical_coord_names = template._logical_coord_names + to_add
+
+        # Patch _add_coordinates to skip adding non-binned dims as 1D dimension coordinates
+        # This prevents them from being added with wrong dimensions (e.g., just "trace")
+        # They will be added later by build_dataset with full spatial_dimension_names
+        _patch_add_coordinates_for_non_binned(template, set(non_binned_dims))
 
 
 def _scan_for_headers(
     segy_file_kwargs: SegyFileArguments,
     segy_file_info: SegyFileInfo,
     template: AbstractDatasetTemplate,
-    grid_overrides: dict[str, Any] | None = None,
+    grid_overrides: GridOverrides | None = None,
 ) -> tuple[list[Dimension], SegyHeaderArray]:
     """Extract trace dimensions and index headers from the SEG-Y file.
 
@@ -346,13 +346,34 @@ def determine_target_size(var_type: str) -> int:
     ds.variables[index].metadata.chunk_grid = chunk_grid
 
 
+def _coerce_grid_overrides(
+    grid_overrides: GridOverrides | dict[str, Any] | None,
+) -> GridOverrides | None:
+    """Normalize public ``grid_overrides`` input into a :class:`GridOverrides` model.
+
+    The internal ingestion pipeline only accepts the typed model. A legacy ``dict`` is
+    converted via :meth:`GridOverrides.from_legacy_dict` and a deprecation message is logged.
+    """
+    if grid_overrides is None:
+        return None
+
+    if isinstance(grid_overrides, GridOverrides):
+        return grid_overrides
+
+    logger.warning(
+        "Passing `grid_overrides` as a dict is deprecated and will be removed in a "
+        "future release; pass a `mdio.GridOverrides` instance instead."
+    )
+    return GridOverrides.model_validate(grid_overrides)
+
+
 def segy_to_mdio(  # noqa PLR0913
     segy_spec: SegySpec,
     mdio_template: AbstractDatasetTemplate,
     input_path: UPath | Path | str,
     output_path: UPath | Path | str,
     overwrite: bool = False,
-    grid_overrides: dict[str, Any] | None = None,
+    grid_overrides: GridOverrides | dict[str, Any] | None = None,
     segy_header_overrides: SegyHeaderOverrides | None = None,
 ) -> None:
     """A function that converts a SEG-Y file to an MDIO v1 file.
@@ -365,12 +386,15 @@ def segy_to_mdio(  # noqa PLR0913
         input_path: The universal path of the input SEG-Y file.
         output_path: The universal path for the output MDIO v1 file.
         overwrite: Whether to overwrite the output file if it already exists. Defaults to False.
-        grid_overrides: Option to add grid overrides.
+        grid_overrides: Option to add grid overrides. Prefer a :class:`mdio.GridOverrides`
+            instance; ``dict`` is still accepted but emits a :class:`DeprecationWarning`.
         segy_header_overrides: Option to override specific SEG-Y headers during ingestion.
 
     Raises:
         FileExistsError: If the output location already exists and overwrite is False.
     """
+    typed_grid_overrides = _coerce_grid_overrides(grid_overrides)
+
     settings = MDIOSettings()
 
     _validate_spec_in_template(segy_spec, mdio_template)
@@ -395,7 +419,7 @@ def segy_to_mdio(  # noqa PLR0913
         segy_file_kwargs,
         segy_file_info,
         template=mdio_template,
-        grid_overrides=grid_overrides,
+        grid_overrides=typed_grid_overrides,
     )
     grid = _build_and_check_grid(segy_dimensions, segy_file_info, segy_headers)
 
@@ -417,7 +441,7 @@ def segy_to_mdio(  # noqa PLR0913
     mdio_template = _update_template_units(mdio_template, spatial_unit)
     mdio_ds: Dataset = mdio_template.build_dataset(name=mdio_template.name, sizes=grid.shape, header_dtype=header_dtype)
 
-    _add_grid_override_to_metadata(dataset=mdio_ds, grid_overrides=grid_overrides)
+    _add_grid_override_to_metadata(dataset=mdio_ds, grid_overrides=typed_grid_overrides)
 
     # Dynamically chunk the variables based on their type
     _chunk_variable(ds=mdio_ds, target_variable_name="trace_mask")  # trace_mask is a Variable and not a Coordinate

src/mdio/ingestion/metadata.py

@@ -3,16 +3,16 @@
 from __future__ import annotations
 
 from typing import TYPE_CHECKING
-from typing import Any
 
 if TYPE_CHECKING:
     from mdio.builder.schemas import Dataset
+    from mdio.segy.geometry import GridOverrides
 
 
-def _add_grid_override_to_metadata(dataset: Dataset, grid_overrides: dict[str, Any] | None) -> None:
+def _add_grid_override_to_metadata(dataset: Dataset, grid_overrides: GridOverrides | None) -> None:
     """Add grid override to Dataset metadata if needed."""
     if dataset.metadata.attributes is None:
         dataset.metadata.attributes = {}
 
     if grid_overrides is not None:
-        dataset.metadata.attributes["gridOverrides"] = grid_overrides
+        dataset.metadata.attributes["gridOverrides"] = grid_overrides.to_legacy_dict()

src/mdio/segy/geometry.py

@@ -6,9 +6,13 @@
 from abc import ABC
 from abc import abstractmethod
 from typing import TYPE_CHECKING
+from typing import Any
 
 import numpy as np
 from numpy.lib import recfunctions as rfn
+from pydantic import BaseModel
+from pydantic import ConfigDict
+from pydantic import Field
 
 from mdio.ingestion.segy.header_analysis import ShotGunGeometryType
 from mdio.ingestion.segy.header_analysis import StreamerShotGeometryType
@@ -31,6 +35,61 @@
 logger = logging.getLogger(__name__)
 
 
+class GridOverrides(BaseModel):
+    """Type-safe configuration for grid override operations during SEG-Y ingestion."""
+
+    model_config = ConfigDict(extra="forbid", validate_by_name=True)
+
+    auto_channel_wrap: bool = Field(
+        default=False,
+        alias="AutoChannelWrap",
+        description="Streamer: auto-detect channel-wrap geometry (Type A vs B).",
+    )
+    auto_shot_wrap: bool = Field(
+        default=False,
+        alias="AutoShotWrap",
+        description="Streamer: derive dense shot_index from interleaved shot_point values.",
+    )
+    calculate_shot_index: bool = Field(
+        default=False,
+        alias="CalculateShotIndex",
+        description="OBN: derive dense shot_index from sparse shot_point values per shot_line.",
+    )
+    non_binned: bool = Field(
+        default=False,
+        alias="NonBinned",
+        description="Collapse selected dims into a single trace dimension without spatial binning.",
+    )
+    has_duplicates: bool = Field(
+        default=False,
+        alias="HasDuplicates",
+        description="Add a trace dimension (chunksize 1) to disambiguate duplicate trace indices.",
+    )
+    chunksize: int | None = Field(
+        default=None,
+        gt=0,
+        description="Chunk size for the trace dimension when `non_binned` is True.",
+    )
+    non_binned_dims: list[str] | None = Field(
+        default=None,
+        description="Dimension names to collapse into the trace dimension when `non_binned` is True.",
+    )
+
+    def __bool__(self) -> bool:
+        """Return True if any override flag is enabled."""
+        return (
+            self.auto_channel_wrap
+            or self.auto_shot_wrap
+            or self.calculate_shot_index
+            or self.non_binned
+            or self.has_duplicates
+        )
+
+    def to_legacy_dict(self) -> dict[str, Any]:
+        """Dump to the legacy ``CamelCase`` dict shape consumed by :class:`GridOverrider`."""
+        return self.model_dump(by_alias=True, exclude_defaults=True)
+
+
 class GridOverrideCommand(ABC):
     """Abstract base class for grid override commands."""
 

src/mdio/segy/utilities.py

@@ -5,7 +5,6 @@
 import itertools
 import logging
 from typing import TYPE_CHECKING
-from typing import Any
 
 import numpy as np
 from dask.array.core import normalize_chunks
@@ -24,6 +23,7 @@
     from mdio.builder.templates.base import AbstractDatasetTemplate
     from mdio.segy.file import SegyFileArguments
     from mdio.segy.file import SegyFileInfo
+    from mdio.segy.geometry import GridOverrides
 
 logger = logging.getLogger(__name__)
 
@@ -34,7 +34,7 @@ def get_grid_plan(  # noqa:  C901, PLR0912, PLR0913, PLR0915
     chunksize: tuple[int, ...] | None,
     template: AbstractDatasetTemplate,
     return_headers: bool = False,
-    grid_overrides: dict[str, Any] | None = None,
+    grid_overrides: GridOverrides | None = None,
 ) -> tuple[list[Dimension], tuple[int, ...]] | tuple[list[Dimension], tuple[int, ...], HeaderArray]:
     """Infer dimension ranges, and increments.
 
@@ -50,17 +50,14 @@ def get_grid_plan(  # noqa:  C901, PLR0912, PLR0913, PLR0915
         chunksize:  Chunk sizes to be used in grid plan.
         template: MDIO template where coordinate names and domain will be taken.
         return_headers: Option to return parsed headers with `Dimension` objects. Default is False.
-        grid_overrides: Option to add grid overrides. See main documentation.
+        grid_overrides: Typed grid override configuration, or ``None`` for no overrides.
 
     Returns:
         All index dimensions and chunksize or dimensions and chunksize together with header values.
 
     Raises:
         ValueError: If computed fields are not found after grid overrides.
     """
-    if grid_overrides is None:
-        grid_overrides = {}
-
     # Keep only dimension and non-dimension coordinates excluding the vertical axis
     horizontal_dimensions = template.spatial_dimension_names
     horizontal_coordinates = horizontal_dimensions + template.coordinate_names
@@ -72,8 +69,8 @@ def get_grid_plan(  # noqa:  C901, PLR0912, PLR0913, PLR0915
     horizontal_coordinates = tuple(c for c in horizontal_coordinates if c not in computed_fields)
 
     # Ensure non_binned_dims are included in the headers to parse, even if not in template
-    if grid_overrides and "non_binned_dims" in grid_overrides:
-        for dim in grid_overrides["non_binned_dims"]:
+    if grid_overrides is not None and grid_overrides.non_binned_dims:
+        for dim in grid_overrides.non_binned_dims:
             if dim not in horizontal_coordinates:
                 horizontal_coordinates = horizontal_coordinates + (dim,)
 
@@ -94,20 +91,24 @@ def get_grid_plan(  # noqa:  C901, PLR0912, PLR0913, PLR0915
         subset=tuple(c for c in horizontal_coordinates if c not in fields_to_skip),
     )
 
-    # Handle grid overrides.
+    # The legacy GridOverrider still consumes the dict shape; dump only at this boundary.
+    # Future PR will replace GridOverrider.
     override_handler = GridOverrider()
     headers_subset, horizontal_coordinates, chunksize = override_handler.run(
         headers_subset,
         horizontal_coordinates,
         chunksize=chunksize,
-        grid_overrides=grid_overrides,
+        grid_overrides=grid_overrides.to_legacy_dict() if grid_overrides is not None else {},
         template=template,
     )
 
     # After grid overrides, determine final spatial dimensions and their chunk sizes
-    non_binned_dims = set()
-    if "NonBinned" in grid_overrides and "non_binned_dims" in grid_overrides:
-        non_binned_dims = set(grid_overrides["non_binned_dims"])
+    non_binned_active = grid_overrides is not None and grid_overrides.non_binned
+    non_binned_dims: set[str] = (
+        set(grid_overrides.non_binned_dims)
+        if non_binned_active and grid_overrides is not None and grid_overrides.non_binned_dims
+        else set()
+    )
 
     # Create mapping from dimension name to original chunk size for easy lookup
     original_spatial_dims = list(template.spatial_dimension_names)
@@ -121,8 +122,11 @@ def get_grid_plan(  # noqa:  C901, PLR0912, PLR0913, PLR0915
         if name in non_binned_dims:
             continue  # Skip dimensions that became coordinates
         if name == "trace":
-            # Special handling for trace dimension
-            chunk_val = int(grid_overrides.get("chunksize", 1)) if "NonBinned" in grid_overrides else 1
+            chunk_val = (
+                int(grid_overrides.chunksize)
+                if non_binned_active and grid_overrides is not None and grid_overrides.chunksize is not None
+                else 1
+            )
             final_spatial_dims.append(name)
             final_spatial_chunks.append(chunk_val)
         elif name in dim_to_chunk:

tests/unit/ingestion/test_metadata.py

@@ -5,6 +5,7 @@
 from types import SimpleNamespace
 
 from mdio.ingestion.metadata import _add_grid_override_to_metadata
+from mdio.segy.geometry import GridOverrides
 
 
 def _make_dataset(attributes: dict | None) -> SimpleNamespace:
@@ -22,18 +23,23 @@ def test_initializes_attributes_dict_when_none(self) -> None:
         assert dataset.metadata.attributes == {}
 
     def test_adds_grid_overrides_when_provided(self) -> None:
-        """Grid overrides should land under the ``gridOverrides`` key."""
+        """Active grid overrides should serialize under the ``gridOverrides`` key."""
         dataset = _make_dataset(attributes=None)
-        overrides = {"HasDuplicates": True, "chunksize": 4}
+        overrides = GridOverrides(has_duplicates=True, chunksize=4)
         _add_grid_override_to_metadata(dataset, grid_overrides=overrides)
-        assert dataset.metadata.attributes == {"gridOverrides": overrides}
+        assert dataset.metadata.attributes == {
+            "gridOverrides": {"HasDuplicates": True, "chunksize": 4},
+        }
 
     def test_preserves_existing_attributes(self) -> None:
         """Existing attribute keys should be preserved when adding overrides."""
         dataset = _make_dataset(attributes={"existing": "value"})
-        overrides = {"NonBinned": True}
+        overrides = GridOverrides(non_binned=True)
         _add_grid_override_to_metadata(dataset, grid_overrides=overrides)
-        assert dataset.metadata.attributes == {"existing": "value", "gridOverrides": overrides}
+        assert dataset.metadata.attributes == {
+            "existing": "value",
+            "gridOverrides": {"NonBinned": True},
+        }
 
     def test_no_overrides_leaves_attributes_untouched(self) -> None:
         """Passing ``None`` overrides must not introduce a ``gridOverrides`` key."""