Linting and cleanup

Author: BrianMichell

src/mdio/core/v1/__init__.py

@@ -4,17 +4,16 @@
 """
 
 from ._overloads import mdio
+from ._serializer import make_coordinate
+from ._serializer import make_dataset
+from ._serializer import make_dataset_metadata
+from ._serializer import make_named_dimension
+from ._serializer import make_variable
 from .builder import MDIODatasetBuilder
 from .builder import write_mdio_metadata
-from ._serializer import (
-    make_coordinate,
-    make_dataset,
-    make_dataset_metadata,
-    make_named_dimension,
-    make_variable,
-)
-from .factory import MDIOSchemaType
 from .factory import SCHEMA_TEMPLATE_MAP
+from .factory import MDIOSchemaType
+
 
 __all__ = [
     "MDIODatasetBuilder",

src/mdio/core/v1/_serializer.py

@@ -4,11 +4,14 @@
 to Zarr storage. This API is not considered stable and may change without notice.
 """
 
-from typing import Any, Dict, List, Optional
+from datetime import datetime
+from typing import Any
+from typing import Dict
+from typing import List
+from typing import Optional
 
 import numpy as np
 from numcodecs import Blosc as NumcodecsBlosc
-from datetime import datetime
 
 from mdio.core.v1._overloads import mdio
 from mdio.schema.compressors import ZFP

src/mdio/core/v1/builder.py

@@ -11,27 +11,27 @@
 
 from zarr.core.chunk_key_encodings import V2ChunkKeyEncoding  # noqa: F401
 
+from mdio.core.v1._overloads import mdio
 from mdio.schema.compressors import ZFP
 from mdio.schema.compressors import Blosc
 from mdio.schema.dimension import NamedDimension
 from mdio.schema.dtype import ScalarType
 from mdio.schema.dtype import StructuredType
 from mdio.schema.metadata import UserAttributes
 from mdio.schema.v1.dataset import Dataset
+from mdio.schema.v1.units import AllUnits
+from mdio.schema.v1.variable import Coordinate
+from mdio.schema.v1.variable import Variable
+from mdio.schema.v1.variable import VariableMetadata
 
 # Import factory functions from serializer module
+from ._serializer import _construct_mdio_dataset
+from ._serializer import _convert_compressor
 from ._serializer import make_coordinate
 from ._serializer import make_dataset
 from ._serializer import make_dataset_metadata
 from ._serializer import make_named_dimension
 from ._serializer import make_variable
-from ._serializer import _convert_compressor
-from ._serializer import _construct_mdio_dataset
-from mdio.core.v1._overloads import mdio
-from mdio.schema.v1.units import AllUnits
-from mdio.schema.v1.variable import Coordinate
-from mdio.schema.v1.variable import Variable
-from mdio.schema.v1.variable import VariableMetadata
 
 
 class _BuilderState(Enum):
@@ -227,9 +227,7 @@ 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: Any) -> 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 +236,7 @@ def write_mdio_metadata(
     Args:
         mdio_ds: The MDIO dataset to serialize
         store: Path to the Zarr store
-        **kwargs: Additional arguments to pass to to_zarr()
+        kwargs: Additional arguments to pass to to_mdio()
 
     Returns:
         The constructed xarray Dataset with MDIO extensions

src/mdio/core/v1/factory.py

@@ -1,24 +1,18 @@
 """MDIO factories for seismic data."""
 
+# TODO(BrianMichell): Add implementations for other canonical datasets.
+
 from __future__ import annotations
 
-import importlib
-from datetime import UTC
-from datetime import datetime
 from enum import Enum
 from enum import auto
 from typing import Any
-from typing import Dict
-from typing import List
-from typing import Optional
 
 from mdio.core.v1.builder import MDIODatasetBuilder
 from mdio.schema.compressors import Blosc
 from mdio.schema.dtype import ScalarType
 from mdio.schema.dtype import StructuredType
 from mdio.schema.v1.dataset import Dataset
-from mdio.schema.v1.units import AllUnits
-from mdio.schema.v1.units import LengthUnitModel
 
 
 class MDIOSchemaType(Enum):
@@ -46,14 +40,14 @@ def __init__(self):
     def create(
         self,
         name: str,
-        shape: List[int],
-        header_fields: Dict[str, str],
+        shape: list[int],
+        header_fields: dict[str, str],
         create_coords: bool = False,
-        sample_format: Optional[str] = None,
-        chunks: Optional[List[int]] = None,
-        sample_units: Optional[Dict[str, str]] = None,
-        z_units: Optional[Dict[str, str]] = None,
-        attributes: Optional[Dict[str, Any]] = None,
+        sample_format: str | None = None,
+        chunks: list[int] | None = None,
+        sample_units: dict[str, str] | None = None,
+        z_units: dict[str, str] | None = None,
+        attributes: dict[str, Any] | None = None,
     ) -> Dataset:
         """Create a generic seismic dataset schema.
 
@@ -80,7 +74,7 @@ def create(
         )
 
         # Add dimensions
-        for dim_name, dim_size in zip(self._dim_names, shape):
+        for dim_name, dim_size in zip(self._dim_names, shape, strict=True):
             builder.add_dimension(
                 name=dim_name,
                 size=dim_size,
@@ -108,10 +102,12 @@ def create(
         )
 
         # Add header variable with structured dtype
-        header_dtype = StructuredType(fields=[
-            {"name": field_name, "format": field_type}
-            for field_name, field_type in header_fields.items()
-        ])
+        header_dtype = StructuredType(
+            fields=[
+                {"name": field_name, "format": field_type}
+                for field_name, field_type in header_fields.items()
+            ]
+        )
         builder.add_variable(
             name="headers",
             data_type=header_dtype,
@@ -145,14 +141,14 @@ def __init__(self, domain: str):
     def create(
         self,
         name: str,
-        shape: List[int],
-        header_fields: Dict[str, str],
+        shape: list[int],
+        header_fields: dict[str, str],
         create_coords: bool = False,
-        sample_format: Optional[str] = None,
-        chunks: Optional[List[int]] = None,
-        sample_units: Optional[Dict[str, str]] = None,
-        z_units: Optional[Dict[str, str]] = None,
-        attributes: Optional[Dict[str, Any]] = None,
+        sample_format: str | None = None,
+        chunks: list[int] | None = None,
+        sample_units: dict[str, str] | None = None,
+        z_units: dict[str, str] | None = None,
+        attributes: dict[str, Any] | None = None,
     ) -> Dataset:
         """Create a seismic dataset schema with domain-specific attributes."""
         # Add seismic-specific attributes
@@ -197,14 +193,14 @@ def __init__(self, domain: str):
     def create(
         self,
         name: str,
-        shape: List[int],
-        header_fields: Dict[str, str],
+        shape: list[int],
+        header_fields: dict[str, str],
         create_coords: bool = False,
-        sample_format: Optional[str] = None,
-        chunks: Optional[List[int]] = None,
-        sample_units: Optional[Dict[str, str]] = None,
-        z_units: Optional[Dict[str, str]] = None,
-        attributes: Optional[Dict[str, Any]] = None,
+        sample_format: str | None = None,
+        chunks: list[int] | None = None,
+        sample_units: dict[str, str] | None = None,
+        z_units: dict[str, str] | None = None,
+        attributes: dict[str, Any] | None = None,
     ) -> Dataset:
         """Create a seismic dataset schema with pre-stack attributes."""
         # Add seismic-specific attributes

tests/integration/test_v1_serialization.py

@@ -4,11 +4,11 @@
 
 import numpy as np
 
-from mdio.core.v1.builder import write_mdio_metadata
 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
 from mdio.core.v1._serializer import make_variable
+from mdio.core.v1.builder import write_mdio_metadata
 from mdio.schema.compressors import ZFP
 from mdio.schema.compressors import Blosc
 from mdio.schema.dtype import ScalarType

tests/unit/test_template_factory.py

@@ -1,88 +1,24 @@
 """Unit tests for MDIO v1 factory."""
 
-# TODO(BrianMichell): Update this to use canonical factory functions. 
+# TODO(BrianMichell): Update this to use canonical factory functions.
 
 from datetime import datetime
 from datetime import timezone
 
-import numpy as np
 import pytest
-import xarray as xr
 from pydantic import ValidationError
-from zarr import Array
 
 from mdio.core.v1._serializer import make_coordinate
 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
 from mdio.core.v1._serializer import make_variable
-from mdio.core.v1.factory import MDIOSchemaType
+from mdio.core.v1.builder import write_mdio_metadata
 from mdio.core.v1.factory import SCHEMA_TEMPLATE_MAP
-from mdio.schema.compressors import ZFP
+from mdio.core.v1.factory import MDIOSchemaType
 from mdio.schema.compressors import Blosc
-from mdio.schema.dimension import NamedDimension
 from mdio.schema.dtype import ScalarType
 from mdio.schema.dtype import StructuredType
-from mdio.schema.v1.dataset import Dataset
-from mdio.core.v1.builder import write_mdio_metadata
-
-# def test_make_named_dimension():
-#     """Test that make_named_dimension returns a NamedDimension object."""
-#     dim = make_named_dimension("time", 42)
-#     assert isinstance(dim, NamedDimension)
-#     assert dim.name == "time"
-#     assert dim.size == 42
-
-
-# def test_make_coordinate_minimal():
-#     """Test that make_coordinate returns a Coordinate object."""
-#     dims = ["x"]
-#     coord = make_coordinate(name="x", dimensions=dims, data_type=ScalarType.FLOAT32)
-#     assert coord.name == "x"
-#     assert coord.dimensions == dims
-#     assert coord.data_type == ScalarType.FLOAT32
-#     assert coord.metadata is None
-
-
-# def test_make_variable_minimal():
-#     """Test that make_variable returns a Variable object."""
-#     var = make_variable(
-#         name="var",
-#         dimensions=["x"],
-#         data_type=ScalarType.FLOAT32,
-#         compressor=None,
-#     )
-#     assert var.name == "var"
-#     assert var.dimensions == ["x"]
-#     assert var.data_type == ScalarType.FLOAT32
-#     assert var.compressor is None
-#     assert var.coordinates is None
-#     assert var.metadata is None
-
-
-# def test_make_dataset_metadata_minimal():
-#     """Test that make_dataset_metadata returns a DatasetMetadata object."""
-#     ts = datetime.now(timezone.utc)
-#     meta = make_dataset_metadata(name="ds", api_version="1", created_on=ts)
-#     assert meta.name == "ds"
-#     assert meta.api_version == "1"
-#     assert meta.created_on == ts
-#     assert meta.attributes is None
-
-
-# def test_make_dataset_minimal():
-#     """Test that make_dataset returns a Dataset object."""
-#     var = make_variable(
-#         name="var",
-#         dimensions=["x"],
-#         data_type=ScalarType.FLOAT32,
-#         compressor=None,
-#     )
-#     ts = datetime.now(timezone.utc)
-#     meta = make_dataset_metadata(name="ds", api_version="1", created_on=ts)
-#     ds = make_dataset([var], meta)
-#     assert ds.variables == [var]
-#     assert ds.metadata == meta
 
 
 def test_make_toy_dataset():