Linting

Author: BrianMichell

src/mdio/core/v1/_overloads.py

@@ -5,6 +5,8 @@
 2. To simplify the API for users where it makes sense (e.g. MDIO v1 uses Zarr and not HDF5).
 """
 
+from collections.abc import Mapping
+
 import xarray as xr
 from xarray import DataArray as _DataArray
 from xarray import Dataset as _Dataset
@@ -15,7 +17,12 @@ class MDIODataset(_Dataset):
 
     __slots__ = ()
 
-    def to_mdio(self, store=None, *args, **kwargs):
+    def to_mdio(
+        self,
+        store: str | None = None,
+        *args: str | int | float | bool,
+        **kwargs: Mapping[str, str | int | float | bool],
+    ) -> None:
         """Alias for `.to_zarr()`, prints a greeting, and writes to Zarr store."""
         print("👋 hello world from mdio.to_mdio!")
         return super().to_zarr(*args, store=store, **kwargs)
@@ -26,7 +33,12 @@ class MDIODataArray(_DataArray):
 
     __slots__ = ()
 
-    def to_mdio(self, store=None, *args, **kwargs):
+    def to_mdio(
+        self,
+        store: str | None = None,
+        *args: str | int | float | bool,
+        **kwargs: Mapping[str, str | int | float | bool],
+    ) -> None:
         """Alias for `.to_zarr()`, prints a greeting, and writes to Zarr store."""
         print("👋 hello world from mdio.to_mdio!")
         return super().to_zarr(*args, store=store, **kwargs)
@@ -39,7 +51,13 @@ class MDIO:
     DataArray = MDIODataArray
 
     @staticmethod
-    def open(store, *args, engine="zarr", consolidated=False, **kwargs):
+    def open(
+        store: str,
+        *args: str | int | float | bool,
+        engine: str = "zarr",
+        consolidated: bool = False,
+        **kwargs: Mapping[str, str | int | float | bool],
+    ) -> MDIODataset:
         """Open a Zarr store as an MDIODataset.
 
         Casts the returned xarray.Dataset (and its variables) to the MDIO subclasses.
@@ -55,9 +73,9 @@ def open(store, *args, engine="zarr", consolidated=False, **kwargs):
         # Cast Dataset to MDIODataset
         ds.__class__ = MDIODataset
         # Cast each DataArray in data_vars and coords
-        for _name, var in ds.data_vars.items():
+        for _name, var in ds.data_vars.values():
             var.__class__ = MDIODataArray
-        for _name, coord in ds.coords.items():
+        for _name, coord in ds.coords.values():
             coord.__class__ = MDIODataArray
         return ds
 

src/mdio/core/v1/_serializer.py

@@ -55,7 +55,7 @@ def make_coordinate(
     )
 
 
-def make_variable(  # noqa: C901
+def make_variable(  # noqa: PLR0913 PLR0912
     name: str,
     dimensions: list[NamedDimension | str],
     data_type: ScalarType | StructuredType,
@@ -118,7 +118,8 @@ def make_variable(  # noqa: C901
         elif isinstance(metadata, VariableMetadata):
             var_metadata = metadata
         else:
-            raise TypeError(f"Unsupported metadata type: {type(metadata)}")
+            msg = f"Unsupported metadata type: {type(metadata)}"
+            raise TypeError(msg)
 
     # Create the variable with all attributes explicitly set
     return Variable(
@@ -170,7 +171,8 @@ def _convert_compressor(
         )
     if isinstance(model, ZFP):
         if zfpy_base is None or ZFPY is None:
-            raise ImportError("zfpy and numcodecs are required to use ZFP compression")
+            msg = "zfpy and numcodecs are required to use ZFP compression"
+            raise ImportError(msg)
         return ZFPY(
             mode=model.mode.value,
             tolerance=model.tolerance,
@@ -179,10 +181,11 @@ def _convert_compressor(
         )
     if model is None:
         return None
-    raise TypeError(f"Unsupported compressor model: {type(model)}")
+    msg = f"Unsupported compressor model: {type(model)}"
+    raise TypeError(msg)
 
 
-def _construct_mdio_dataset(mdio_ds: MDIODataset) -> mdio.Dataset:  # noqa: C901
+def _construct_mdio_dataset(mdio_ds: MDIODataset) -> mdio.Dataset:  # noqa: PLR0912
     """Build an MDIO dataset with correct dimensions and dtypes.
 
     This internal function constructs the underlying data structure for an MDIO dataset,
@@ -215,7 +218,8 @@ def _construct_mdio_dataset(mdio_ds: MDIODataset) -> mdio.Dataset:  # noqa: C901
         elif isinstance(dt, StructuredType):
             dtype = np.dtype([(f.name, f.format.value) for f in dt.fields])
         else:
-            raise TypeError(f"Unsupported data_type: {dt}")
+            msg = f"Unsupported data_type: {dt}"
+            raise TypeError(msg)
         arr = np.zeros(shape, dtype=dtype)
         data_array = mdio.DataArray(arr, dims=dim_names)
         data_array.encoding["fill_value"] = 0.0

src/mdio/core/v1/builder.py

@@ -1,5 +1,6 @@
 """Builder pattern implementation for MDIO v1 schema models."""
 
+from collections.abc import Mapping
 from datetime import UTC
 from datetime import datetime
 from enum import Enum
@@ -53,7 +54,7 @@ class MDIODatasetBuilder:
 
     def __init__(self, name: str, attributes: dict[str, Any] | None = None):
         self.name = name
-        self.api_version = "1.0.0"  # TODO: Pull from package metadata
+        self.api_version = "1.0.0"  # TODO(BrianMichell, #0): Pull from package metadata
         self.created_on = datetime.now(UTC)
         self.attributes = attributes
         self._dimensions: list[NamedDimension] = []
@@ -62,7 +63,7 @@ def __init__(self, name: str, attributes: dict[str, Any] | None = None):
         self._state = _BuilderState.INITIAL
         self._unnamed_variable_counter = 0
 
-    def add_dimension(
+    def add_dimension(  # noqa: PLR0913
         self,
         name: str,
         size: int,
@@ -101,7 +102,7 @@ def add_dimension(
         self._state = _BuilderState.HAS_DIMENSIONS
         return self
 
-    def add_coordinate(
+    def add_coordinate(  # noqa: PLR0913
         self,
         name: str = "",
         *,
@@ -112,7 +113,8 @@ def add_coordinate(
     ) -> "MDIODatasetBuilder":
         """Add a coordinate after adding at least one dimension."""
         if self._state == _BuilderState.INITIAL:
-            raise ValueError("Must add at least one dimension before adding coordinates")
+            msg = "Must add at least one dimension before adding coordinates"
+            raise ValueError(msg)
 
         if name == "":
             name = f"coord_{len(self._coordinates)}"
@@ -127,7 +129,8 @@ def add_coordinate(
             if isinstance(dim, str):
                 dim_obj = next((d for d in self._dimensions if d.name == dim), None)
                 if dim_obj is None:
-                    raise ValueError(f"Dimension {dim!r} not found")
+                    msg = f"Dimension {dim!r} not found"
+                    raise ValueError(msg)
                 dim_objects.append(dim_obj)
             else:
                 dim_objects.append(dim)
@@ -144,7 +147,7 @@ def add_coordinate(
         self._state = _BuilderState.HAS_COORDINATES
         return self
 
-    def add_variable(
+    def add_variable(  # noqa: PLR0913
         self,
         name: str = "",
         *,
@@ -157,7 +160,8 @@ def add_variable(
     ) -> "MDIODatasetBuilder":
         """Add a variable after adding at least one dimension."""
         if self._state == _BuilderState.INITIAL:
-            raise ValueError("Must add at least one dimension before adding variables")
+            msg = "Must add at least one dimension before adding variables"
+            raise ValueError(msg)
 
         if name == "":
             name = f"var_{self._unnamed_variable_counter}"
@@ -171,7 +175,8 @@ def add_variable(
             if isinstance(dim, str):
                 dim_obj = next((d for d in self._dimensions if d.name == dim), None)
                 if dim_obj is None:
-                    raise ValueError(f"Dimension {dim!r} not found")
+                    msg = f"Dimension {dim!r} not found"
+                    raise ValueError(msg)
                 dim_objects.append(dim_obj)
             else:
                 dim_objects.append(dim)
@@ -193,7 +198,8 @@ def add_variable(
     def build(self) -> Dataset:
         """Build the final dataset."""
         if self._state == _BuilderState.INITIAL:
-            raise ValueError("Must add at least one dimension before building")
+            msg = "Must add at least one dimension before building"
+            raise ValueError(msg)
 
         metadata = make_dataset_metadata(
             self.name, self.api_version, self.created_on, self.attributes
@@ -216,7 +222,11 @@ def build(self) -> Dataset:
         return make_dataset(all_variables, metadata)
 
 
-def write_mdio_metadata(mdio_ds: Dataset, store: str, **kwargs: Any) -> mdio.Dataset:
+def write_mdio_metadata(
+    mdio_ds: Dataset,
+    store: str,
+    **kwargs: Mapping[str, str | int | float | bool],
+) -> mdio.Dataset:
     """Write MDIO metadata to a Zarr store and return the constructed mdio.Dataset.
 
     This function constructs an mdio.Dataset from the MDIO dataset and writes its metadata
@@ -238,7 +248,7 @@ def _generate_encodings() -> dict:
         Returns:
             Dictionary mapping variable names to their encoding configurations.
         """
-        # TODO: Re-enable chunk_key_encoding when supported by xarray
+        # TODO(Anybody, #10274): Re-enable chunk_key_encoding when supported by xarray
         # dimension_separator_encoding = V2ChunkKeyEncoding(separator="/").to_dict()
         global_encodings = {}
         for var in mdio_ds.variables:
@@ -250,7 +260,7 @@ def _generate_encodings() -> dict:
                 chunks = var.metadata.chunk_grid.configuration.chunk_shape
             global_encodings[var.name] = {
                 "chunks": chunks,
-                # TODO: Re-enable chunk_key_encoding when supported by xarray
+                # TODO(Anybody, #10274): Re-enable chunk_key_encoding when supported by xarray
                 # "chunk_key_encoding": dimension_separator_encoding,
                 "_FillValue": fill_value,
                 "dtype": var.data_type,

src/mdio/core/v1/factory.py

@@ -1,18 +1,21 @@
 """MDIO factories for seismic data."""
 
-# TODO(BrianMichell): Add implementations for other canonical datasets.
+# TODO(BrianMichell, #535): Add implementations for other canonical datasets.
 
 from __future__ import annotations
 
 from enum import Enum
 from enum import auto
+from typing import TYPE_CHECKING
 from typing import Any
 
 from mdio.core.v1.builder import MDIODatasetBuilder
 from mdio.schemas.compressors import Blosc
 from mdio.schemas.dtype import ScalarType
 from mdio.schemas.dtype import StructuredType
-from mdio.schemas.v1.dataset import Dataset
+
+if TYPE_CHECKING:
+    from mdio.schemas.v1.dataset import Dataset
 
 
 class MDIOSchemaType(Enum):
@@ -28,15 +31,15 @@ class MDIOSchemaType(Enum):
 class Seismic3DPostStackGeneric:
     """Generic 3D seismic post stack dataset."""
 
-    def __init__(self):
+    def __init__(self) -> None:
         self._dim_names = ["inline", "crossline", "sample"]
         self._chunks = [128, 128, 128]  # 8 mb
         self._coords = {
             "cdp-x": ("float64", {"unitsV1": {"length": "m"}}, self._dim_names[:-1]),
             "cdp-y": ("float64", {"unitsV1": {"length": "m"}}, self._dim_names[:-1]),
         }
 
-    def create(
+    def create(  # noqa: PLR0913
         self,
         name: str,
         shape: list[int],
@@ -128,11 +131,11 @@ def create(
 class Seismic3DPostStack(Seismic3DPostStackGeneric):
     """3D seismic post stack dataset with domain-specific attributes."""
 
-    def __init__(self, domain: str):
+    def __init__(self, domain: str) -> None:
         super().__init__()
         self._dim_names = ["inline", "crossline", domain]
 
-    def create(
+    def create(  # noqa: PLR0913
         self,
         name: str,
         shape: list[int],
@@ -170,7 +173,7 @@ def create(
 class Seismic3DPreStack(Seismic3DPostStackGeneric):
     """3D seismic pre stack dataset."""
 
-    def __init__(self, domain: str):
+    def __init__(self, domain: str) -> None:
         super().__init__()
         self._dim_names = ["inline", "crossline", "offset", domain]
         self._chunks = [1, 1, 512, 4096]  # 8 mb
@@ -179,7 +182,7 @@ def __init__(self, domain: str):
             "cdp-y": ("float64", {"length": "m"}, self._dim_names[:-2]),
         }
 
-    def create(
+    def create(  # noqa: PLR0913
         self,
         name: str,
         shape: list[int],

tests/integration/test_v1_serialization.py

@@ -1,9 +1,11 @@
 """Integration test for MDIO v1 Xarray Zarr constructor."""
 
 from datetime import datetime
+from pathlib import Path
 
 import numpy as np
 
+from mdio.core.v1._overloads import MDIODataset
 from mdio.core.v1._serializer import make_dataset
 from mdio.core.v1._serializer import make_dataset_metadata
 from mdio.core.v1._serializer import make_named_dimension
@@ -15,7 +17,7 @@
 from mdio.schemas.dtype import StructuredType
 
 
-def build_toy_dataset():
+def build_toy_dataset() -> MDIODataset:
     """Build a toy dataset for testing."""
     # core dimensions
     inline = make_named_dimension("inline", 256)
@@ -165,7 +167,7 @@ def build_toy_dataset():
     )
 
 
-def test_to_mdio_writes_and_returns_mdio(tmp_path):
+def test_to_mdio_writes_and_returns_mdio(tmp_path: Path) -> None:
     """Test that to_mdio writes and returns an mdio.Dataset."""
     ds_in = build_toy_dataset()
     store_path = tmp_path / "toy.mdio"