Update overloads

Author: BrianMichell

src/mdio/core/v1/__init__.py

@@ -1,27 +1,12 @@
-# mdio/__init__.py
+"""
+MDIO core v1 package initialization.
+Exposes the MDIO overloads and core v1 functionality.
+"""
 
-import xarray as _xr
-from ._overloads import open_mdio, to_mdio
+from ._overloads import mdio
+from .constructor import Write_MDIO_metadata
 
 __all__ = [
-    # explicit overrides / aliases
-    "open_mdio",
-    "to_mdio",
-    # everything else will be auto-populated by __dir__ / __getattr__
+    "mdio",
+    "Write_MDIO_metadata",
 ]
-
-def __getattr__(name: str):
-    """
-    Fallback: anything not defined in mdio/__init__.py
-    gets looked up on xarray.
-    """
-    if hasattr(_xr, name):
-        return getattr(_xr, name)
-    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
-
-def __dir__():
-    """
-    Make dir(mdio) list our overrides and then all public xarray names.
-    """
-    xr_public = [n for n in dir(_xr) if not n.startswith("_")]
-    return sorted(__all__ + xr_public)

src/mdio/core/v1/_overloads.py

@@ -5,32 +5,62 @@
 1. To provide a consistent mdio.* naming scheme.
 2. To simplify the API for users where it makes sense (e.g. MDIO v1 uses Zarr and not HDF5).
 """
+import xarray as xr
+from xarray import Dataset as _Dataset, DataArray as _DataArray
 
 
-import xarray as _xr
-from xarray import Dataset as _Dataset, DataArray as _DataArray
+class MDIODataset(_Dataset):
+    """xarray.Dataset subclass with MDIO v1 extensions."""
+    __slots__ = ()
+
+    def to_mdio(self, store=None, *args, **kwargs):
+        """
+        Alias for `.to_zarr()`, prints a greeting, and writes to Zarr store.
+        """
+        print("👋 hello world from mdio.to_mdio!")
+        return super().to_zarr(store=store, *args, **kwargs)
+
+
+class MDIODataArray(_DataArray):
+    """xarray.DataArray subclass with MDIO v1 extensions."""
+    __slots__ = ()
+
+    def to_mdio(self, store=None, *args, **kwargs):
+        """
+        Alias for `.to_zarr()`, prints a greeting, and writes to Zarr store.
+        """
+        print("👋 hello world from mdio.to_mdio!")
+        return super().to_zarr(store=store, *args, **kwargs)
+
+
+class MDIO:
+    """MDIO namespace for overloaded types and functions."""
+    Dataset = MDIODataset
+    DataArray = MDIODataArray
+
+    @staticmethod
+    def open(store, *args, engine="zarr", consolidated=False, **kwargs):
+        """
+        Open a Zarr store as an MDIODataset. Prints a greeting and casts
+        the returned xarray.Dataset (and its variables) to the MDIO subclasses.
+        """
+        print("👋 hello world from mdio.open!")
+        ds = xr.open_dataset(
+            store,
+            engine=engine,
+            consolidated=consolidated,
+            *args,
+            **kwargs,
+        )
+        # Cast Dataset to MDIODataset
+        ds.__class__ = MDIODataset
+        # Cast each DataArray in data_vars and coords
+        for name, var in ds.data_vars.items():
+            var.__class__ = MDIODataArray
+        for name, coord in ds.coords.items():
+            coord.__class__ = MDIODataArray
+        return ds
+
 
-def open_mdio(store, *args, engine="zarr", consolidated=False, **kwargs):
-    """
-    Our mdio version of xr.open_zarr. Prints a greeting,
-    then calls xr.open_dataset(..., engine="zarr").
-    """
-    print("👋 hello world from mdio.open_mdio!")
-    return _xr.open_dataset(store, *args,
-                             engine=engine,
-                             consolidated=consolidated,
-                             **kwargs)
-
-def to_mdio(self, *args, **kwargs):
-    """
-    Alias for .to_zarr, renamed to .to_mdio,
-    so you get a consistent mdio.* naming.
-    """
-    print("👋 hello world from mdio.to_mdio!")
-    print(f"kwargs: {kwargs}")
-    return self.to_zarr(*args, **kwargs)
-
-# Monkey-patch Dataset and DataArray so that you can do:
-#   ds.to_mdio(...)        and     arr.to_mdio(...)
-_Dataset.to_mdio = to_mdio
-_DataArray.to_mdio = to_mdio
+# Create module-level MDIO namespace
+mdio = MDIO()

src/mdio/core/v1/constructor.py

@@ -1,4 +1,4 @@
-"""Construct an Xarray Dataset from an MDIO v1 Dataset and write to Zarr."""
+"""Construct an MDIO dataset and write to Zarr."""
 import xarray as xr
 import numpy as np
 import dask.array as da
@@ -10,18 +10,19 @@
 from mdio.schema.dtype import ScalarType, StructuredType
 from mdio.schema.compressors import Blosc, ZFP
 from mdio.schema.v1.variable import Coordinate
-
+from mdio.core.v1._overloads import mdio
 
 from numcodecs import Blosc as NumcodecsBlosc
 
 try:
-    import zfpy as BaseZFPY # Baser library
+    import zfpy as BaseZFPY  # Base library
     from numcodecs import ZFPY as NumcodecsZFPY  # Codec
 except ImportError:
-    print(f"Tried to import zfpy and numcodes zfpy but failed because {ImportError}")
+    print(f"Tried to import zfpy and numcodecs zfpy but failed because {ImportError}")
     BaseZFPY = None
     NumcodecsZFPY = None
 
+
 def _convert_compressor(model: Blosc | ZFP | None) -> NumcodecsBlosc | NumcodecsZFPY | None:
     if isinstance(model, Blosc):
         return NumcodecsBlosc(
@@ -45,7 +46,7 @@ def _convert_compressor(model: Blosc | ZFP | None) -> NumcodecsBlosc | Numcodecs
         raise TypeError(f"Unsupported compressor model: {type(model)}")
 
 
-def _construct_mdio_dataset(mdio_ds: MDIODataset) -> xr.Dataset:
+def _construct_mdio_dataset(mdio_ds: MDIODataset) -> mdio.Dataset:
     """Build an MDIO dataset with correct dimensions and dtypes.
     
     This internal function constructs the underlying data structure for an MDIO dataset,
@@ -65,7 +66,7 @@ def _construct_mdio_dataset(mdio_ds: MDIODataset) -> xr.Dataset:
                 dims[d.name] = d.size
 
     # Build data variables
-    data_vars: dict[str, xr.DataArray] = {}
+    data_vars: dict[str, mdio.DataArray] = {}
     for var in mdio_ds.variables:
         dim_names = [d.name if isinstance(d, NamedDimension) else d for d in var.dimensions]
         shape = tuple(dims[name] for name in dim_names)
@@ -76,12 +77,8 @@ def _construct_mdio_dataset(mdio_ds: MDIODataset) -> xr.Dataset:
             dtype = np.dtype([(f.name, f.format.value) for f in dt.fields])
         else:
             raise TypeError(f"Unsupported data_type: {dt}")
-        # arr = da.zeros(shape, dtype=dtype)
         arr = np.zeros(shape, dtype=dtype)
-        data_array = xr.DataArray(arr, dims=dim_names)
-        # set default fill_value to zero instead of NaN
-        # TODO: This seems to be ignored by xarray.
-        # Setting in the _generate_encodings() function does work though.
+        data_array = mdio.DataArray(arr, dims=dim_names)
         data_array.encoding['fill_value'] = 0.0
 
         # Set long_name if present
@@ -90,33 +87,30 @@ def _construct_mdio_dataset(mdio_ds: MDIODataset) -> xr.Dataset:
 
         # Set coordinates if present, excluding dimension names
         if var.coordinates is not None:
-            # Get the set of dimension names for this variable
             dim_set = set(dim_names)
-            # Filter out any coordinates that are also dimensions
             coord_names = [
-                c.name if isinstance(c, Coordinate) else c 
-                for c in var.coordinates 
+                c.name if isinstance(c, Coordinate) else c  
+                for c in var.coordinates  
                 if (c.name if isinstance(c, Coordinate) else c) not in dim_set
             ]
-            if coord_names:  # Only set coordinates if there are any non-dimension coordinates
+            if coord_names:
                 data_array.attrs["coordinates"] = " ".join(coord_names)
 
-        # attach variable metadata into DataArray attributes, excluding nulls and chunkGrid
+        # Attach variable metadata into DataArray attributes
         if var.metadata is not None:
             md = var.metadata.model_dump(
                 by_alias=True,
                 exclude_none=True,
                 exclude={"chunk_grid"},
             )
-            # Convert single-element lists to objects
             for key, value in md.items():
                 if isinstance(value, list) and len(value) == 1:
                     md[key] = value[0]
             data_array.attrs.update(md)
         data_vars[var.name] = data_array
 
-    ds = xr.Dataset(data_vars)
-    # Attach metadata as attrs
+    ds = mdio.Dataset(data_vars)
+    # Attach dataset metadata
     ds.attrs["apiVersion"] = mdio_ds.metadata.api_version
     ds.attrs["createdOn"] = str(mdio_ds.metadata.created_on)
     ds.attrs["name"] = mdio_ds.metadata.name
@@ -125,16 +119,13 @@ def _construct_mdio_dataset(mdio_ds: MDIODataset) -> xr.Dataset:
     return ds
 
 
-
-
-def Write_MDIO_metadata(mdio_ds: MDIODataset, store: str, **kwargs: Any) -> xr.Dataset:
-    """Write MDIO metadata to a Zarr store and return the constructed xarray.Dataset.
+def Write_MDIO_metadata(mdio_ds: MDIODataset, store: str, **kwargs: Any) -> mdio.Dataset:
+    """Write MDIO metadata to a Zarr store and return the constructed mdio.Dataset.
     
-    This function constructs an xarray.Dataset from the MDIO dataset and writes its metadata
+    This function constructs an mdio.Dataset from the MDIO dataset and writes its metadata
     to a Zarr store. The actual data is not written, only the metadata structure is created.
     """
     ds = _construct_mdio_dataset(mdio_ds)
-    # Write to Zarr format v2 with consolidated metadata and all attributes
 
     def _generate_encodings() -> dict:
         """Generate encodings for each variable in the MDIO dataset.
@@ -147,29 +138,28 @@ def _generate_encodings() -> dict:
         for var in mdio_ds.variables:
             fill_value = 0
             if isinstance(var.data_type, StructuredType):
-                # Create a structured fill value that matches the dtype
-                # fill_value = np.zeros(1, dtype=[(f.name, f.format.value) for f in var.data_type.fields])[0]
-                # TODO: Re-enable this once xarray supports this PR https://github.com/zarr-developers/zarr-python/pull/3015
                 continue
             chunks = None
             if var.metadata is not None and var.metadata.chunk_grid is not None:
                 chunks = var.metadata.chunk_grid.configuration.chunk_shape
             global_encodings[var.name] = {
                 "chunks": chunks,
-                # TODO: Re-enable this once xarray supports this PR https://github.com/pydata/xarray/pull/10274
+                # TODO: Re-enable chunk_key_encoding when supported by xarray
                 # "chunk_key_encoding": dimension_separator_encoding,
                 "_FillValue": fill_value,
                 "dtype": var.data_type,
                 "compressors": _convert_compressor(var.compressor),
             }
         return global_encodings
 
-    ds.to_mdio(store, 
-               mode="w", 
-               zarr_format=2,
-               consolidated=True,
-               safe_chunks=False,  # This ignores the Dask chunks
-               compute=False,  # Ensures only the metadata is written
-               encoding=_generate_encodings(),
-               **kwargs)
-    return ds
+    ds.to_mdio(
+        store,
+        mode="w",
+        zarr_format=2,
+        consolidated=True,
+        safe_chunks=False,
+        compute=False,
+        encoding=_generate_encodings(),
+        **kwargs
+    )
+    return ds