add schemas

Author: tasansal

src/mdio/schemas/__init__.py

@@ -0,0 +1,18 @@
+"""MDIO schemas for different data types."""
+
+from mdio.schemas.compressors import ZFP
+from mdio.schemas.compressors import Blosc
+from mdio.schemas.dimension import NamedDimension
+from mdio.schemas.dtype import ScalarType
+from mdio.schemas.dtype import StructuredField
+from mdio.schemas.dtype import StructuredType
+
+
+__all__ = [
+    "Blosc",
+    "ZFP",
+    "NamedDimension",
+    "ScalarType",
+    "StructuredField",
+    "StructuredType",
+]

src/mdio/schemas/base.py

@@ -0,0 +1,43 @@
+"""Base models to subclass from."""
+
+from pydantic import ConfigDict
+from pydantic import Field
+from pydantic.json_schema import GenerateJsonSchema
+
+from mdio.schemas.compressors import ZFP
+from mdio.schemas.compressors import Blosc
+from mdio.schemas.core import CamelCaseStrictModel
+from mdio.schemas.dimension import NamedDimension
+from mdio.schemas.dtype import DataTypeModel
+
+
+JSON_SCHEMA_DIALECT = GenerateJsonSchema.schema_dialect
+
+
+class BaseDataset(CamelCaseStrictModel):
+    """A base class for MDIO datasets.
+
+    We add schema dialect to extend the config of `StrictCamelBaseModel`.
+    We use the default Pydantic schema generator `GenerateJsonSchema` to
+    define the JSON schema dialect accurately.
+    """
+
+    model_config = ConfigDict(json_schema_extra={"$schema": JSON_SCHEMA_DIALECT})
+
+
+class BaseArray(DataTypeModel, CamelCaseStrictModel):
+    """A base array schema."""
+
+    dimensions: list[NamedDimension] | list[str] = Field(
+        ..., description="List of Dimension collection or reference to dimension names."
+    )
+    compressor: Blosc | ZFP | None = Field(
+        default=None, description="Compression settings."
+    )
+
+
+class NamedArray(BaseArray):
+    """An array with a name."""
+
+    name: str = Field(..., description="Name of the array.")
+    long_name: str | None = Field(default=None, description="Fully descriptive name.")

src/mdio/schemas/builder.py

@@ -0,0 +1,156 @@
+"""Schema builders."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from mdio.schemas import NamedDimension
+from mdio.schemas.v1.dataset import Dataset
+from mdio.schemas.v1.dataset import DatasetMetadata
+from mdio.schemas.v1.variable import Variable
+from mdio.schemas.v1.variable import VariableMetadata
+
+
+class VariableBuilder:
+    """Dataset builder."""
+
+    def __init__(self) -> None:
+        """Initialize the builder."""
+        self.name = None
+        self.long_name = None
+        self.dtype = None
+        self.chunks = None
+        self.dims = None
+        self.coords = None
+        self.compressor = None
+        self.meta_dict = None
+
+    def set_name(self, name: str) -> VariableBuilder:
+        """Set variable name."""
+        self.name = name
+        return self
+
+    def set_long_name(self, long_name: str) -> VariableBuilder:
+        """Add long, descriptive name to the variable."""
+        self.long_name = long_name
+        return self
+
+    def set_compressor(self, compressor: dict[str, Any]) -> VariableBuilder:
+        """Add long, descriptive name to the variable."""
+        self.compressor = compressor
+        return self
+
+    def add_dimension(self, *dimensions: str | dict[str, int]) -> VariableBuilder:
+        """Add a dimension to the dataset."""
+        if self.dims is None:
+            self.dims = []
+
+        if isinstance(dimensions[0], str):
+            dims = list(dimensions)
+        elif isinstance(dimensions[0], dict):
+            dims = [
+                NamedDimension(name=name, size=size)
+                for dim in dimensions
+                for name, size in dim.items()
+            ]
+        else:
+            raise NotImplementedError
+
+        self.dims.extend(dims)
+        return self
+
+    def add_coordinate(self, *names: str) -> VariableBuilder:
+        """Add a coordinate to the variable."""
+        if self.coords is None:
+            self.coords = []
+
+        self.coords.extend(names)
+        return self
+
+    def set_format(self, format_: str | dict[str, str]) -> VariableBuilder:
+        """Set variable format."""
+        if isinstance(format_, dict):
+            fields = [{"name": n, "format": f} for n, f in format_.items()]
+            format_ = {"fields": fields}
+
+        self.dtype = format_
+        return self
+
+    def set_chunks(self, chunks: list[int]) -> VariableBuilder:
+        """Set variable chunks."""
+        if self.meta_dict is None:
+            self.meta_dict = {}
+
+        self.meta_dict["chunkGrid"] = {"configuration": {"chunkShape": chunks}}
+        return self
+
+    def set_units(self, units: dict[str, str]) -> VariableBuilder:
+        """Set variable units."""
+        if self.meta_dict is None:
+            self.meta_dict = {}
+
+        self.meta_dict["unitsV1"] = units
+        return self
+
+    def add_attribute(self, key: str, value: Any) -> VariableBuilder:  # noqa: ANN401
+        """Add a user attribute to the variable metadata."""
+        if self.meta_dict is None:
+            self.meta_dict = {}
+
+        self.meta_dict["attributes"] = {key: value}
+        return self
+
+    def build(self) -> Variable:
+        """Build the dataset model."""
+        if self.chunks is not None and len(self.chunks) != len(self.dims):
+            msg = "Variable chunks must have same number of dimensions."
+            raise ValueError(msg)
+
+        var_kwargs = {}
+
+        if self.meta_dict is not None:
+            var_kwargs["metadata"] = VariableMetadata.model_validate(self.meta_dict)
+
+        return Variable(
+            name=self.name,
+            long_name=self.long_name,
+            data_type=self.dtype,
+            dimensions=self.dims,
+            coordinates=self.coords,
+            compressor=self.compressor,
+            **var_kwargs,
+        )
+
+
+class DatasetBuilder:
+    """Dataset builder."""
+
+    def __init__(self) -> None:
+        """Initialize the builder."""
+        self.variables = []
+        self.name = None
+        self.metadata = None
+
+    def set_name(self, name: str) -> DatasetBuilder:
+        """Set dataset name."""
+        self.name = name
+        return self
+
+    def add_variable(self, variable: Variable) -> DatasetBuilder:
+        """Add a variable to the dataset."""
+        self.variables.append(variable)
+        return self
+
+    def add_variables(self, variables: list[Variable]) -> DatasetBuilder:
+        """Add multiple variables to the dataset."""
+        [self.add_variable(variable) for variable in variables]
+        return self
+
+    def set_metadata(self, metadata: DatasetMetadata) -> DatasetBuilder:
+        """Add a metadata to the dataset."""
+        self.metadata = metadata
+        return self
+
+    def build(self) -> Dataset:
+        """Build the dataset model."""
+        return Dataset(variables=self.variables, metadata=self.metadata)

src/mdio/schemas/chunk_grid.py

@@ -0,0 +1,44 @@
+"""This module contains data models for Zarr's chunk grid."""
+
+from __future__ import annotations
+
+from pydantic import Field
+
+from mdio.schemas.core import CamelCaseStrictModel
+
+
+class RegularChunkShape(CamelCaseStrictModel):
+    """Represents regular chunk sizes along each dimension."""
+
+    chunk_shape: list[int] = Field(
+        ..., description="Lengths of the chunk along each dimension of the array."
+    )
+
+
+class RectilinearChunkShape(CamelCaseStrictModel):
+    """Represents irregular chunk sizes along each dimension."""
+
+    chunk_shape: list[list[int]] = Field(
+        ...,
+        description="Lengths of the chunk along each dimension of the array.",
+    )
+
+
+class RegularChunkGrid(CamelCaseStrictModel):
+    """Represents a rectangular and regularly spaced chunk grid."""
+
+    name: str = Field(default="regular", description="The name of the chunk grid.")
+
+    configuration: RegularChunkShape = Field(
+        ..., description="Configuration of the regular chunk grid."
+    )
+
+
+class RectilinearChunkGrid(CamelCaseStrictModel):
+    """Represents a rectangular and irregularly spaced chunk grid."""
+
+    name: str = Field(default="rectilinear", description="The name of the chunk grid.")
+
+    configuration: RectilinearChunkShape = Field(
+        ..., description="Configuration of the irregular chunk grid."
+    )