[feat] Implemented flex shape support and refactored validators

Author: yhteoh

src/oqd_dataschema/__init__.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .base import Dataset, GroupBase, GroupRegistry, condataset
+from .base import CastDataset, Dataset, GroupBase, GroupRegistry, condataset
 from .datastore import Datastore
 from .groups import (
     ExpectationValueDataGroup,
@@ -24,6 +24,7 @@
 ########################################################################################
 
 __all__ = [
+    "CastDataset",
     "Dataset",
     "Datastore",
     "GroupBase",

src/oqd_dataschema/base.py

@@ -16,7 +16,7 @@
 import typing
 import warnings
 from enum import Enum
-from functools import partial, reduce
+from types import NoneType
 from typing import Annotated, Any, ClassVar, Literal, Optional, Sequence, Tuple, Union
 
 import numpy as np
@@ -28,9 +28,12 @@
     Discriminator,
     Field,
     TypeAdapter,
+    field_validator,
     model_validator,
 )
 
+from .utils import _flex_shape_equal, _validator_from_condition
+
 ########################################################################################
 
 __all__ = ["GroupBase", "Dataset", "GroupRegistry", "condataset", "CastDataset"]
@@ -108,7 +111,7 @@ class Dataset(BaseModel, extra="forbid"):
     """
 
     dtype: Optional[Literal[DTypes.names()]] = None  # type: ignore
-    shape: Optional[Tuple[int, ...]] = None
+    shape: Optional[Tuple[Union[int, None], ...]] = None
     data: Optional[Any] = Field(default=None, exclude=True)
 
     attrs: Attrs = {}
@@ -117,51 +120,58 @@ class Dataset(BaseModel, extra="forbid"):
         use_enum_values=False, arbitrary_types_allowed=True, validate_assignment=True
     )
 
+    @field_validator("data", mode="before")
     @classmethod
-    def cast(cls, data):
-        if isinstance(data, np.ndarray):
-            return cls(data=data)
-        return data
+    def validate_and_update(cls, value):
+        # check if data exist
+        if value is None:
+            return value
 
-    @model_validator(mode="before")
-    @classmethod
-    def validate_and_update(cls, values: dict):
-        data = values.get("data")
-        dtype = values.get("dtype")
-        shape = values.get("shape")
+        # check if data is a numpy array
+        if not isinstance(value, np.ndarray):
+            raise TypeError("`data` must be a numpy.ndarray.")
 
-        if data is None and (dtype is not None and shape is not None):
-            return values
+        return value
 
-        elif data is not None and (dtype is None and shape is None):
-            if not isinstance(data, np.ndarray):
-                raise TypeError("`data` must be a numpy.ndarray.")
+    @model_validator(mode="after")
+    def validate_data_matches_shape_dtype(self):
+        """Ensure that `data` matches `dtype` and `shape`."""
 
-            if type(data.dtype) not in DTypes:
-                raise TypeError(
-                    f"`data` must be a numpy array of dtype in {tuple(DTypes.names())}."
-                )
+        # check if data exist
+        if self.data is None:
+            return self
+
+        # check if dtype matches data
+        if (
+            self.dtype is not None
+            and type(self.data.dtype) is not DTypes.get(self.dtype).value
+        ):
+            raise ValueError(
+                f"Expected data dtype `{self.dtype}`, but got `{self.data.dtype.name}`."
+            )
 
-            values["dtype"] = DTypes(type(data.dtype)).name.lower()
-            values["shape"] = data.shape
+        # check if shape mataches data
+        if self.shape is not None and not _flex_shape_equal(
+            self.data.shape, self.shape
+        ):
+            raise ValueError(f"Expected shape {self.shape}, but got {self.data.shape}.")
 
-        return values
+        # reassign dtype if it is None
+        if self.dtype != DTypes(type(self.data.dtype)).name.lower():
+            self.dtype = DTypes(type(self.data.dtype)).name.lower()
+
+        # resassign shape to concrete value if it is None or a flexible shape
+        if self.shape != self.data.shape:
+            self.shape = self.data.shape
 
-    @model_validator(mode="after")
-    def validate_data_matches_shape_dtype(self):
-        """Ensure that `data` matches `dtype` and `shape`."""
-        if self.data is not None:
-            expected_dtype = DTypes.get(self.dtype).value
-            if type(self.data.dtype) is not expected_dtype:
-                raise ValueError(
-                    f"Expected data dtype `{self.dtype}`, but got `{self.data.dtype.name}`."
-                )
-            if self.data.shape != self.shape:
-                raise ValueError(
-                    f"Expected shape {self.shape}, but got {self.data.shape}."
-                )
         return self
 
+    @classmethod
+    def cast(cls, data):
+        if isinstance(data, np.ndarray):
+            return cls(data=data)
+        return data
+
     def __getitem__(self, idx):
         return self.data[idx]
 
@@ -172,6 +182,12 @@ def _is_dataset_type(cls, type_):
         )
 
 
+CastDataset = Annotated[Dataset, BeforeValidator(Dataset.cast)]
+
+########################################################################################
+
+
+@_validator_from_condition
 def _constrain_dtype(dataset, *, dtype_constraint=None):
     if (not isinstance(dtype_constraint, str)) and isinstance(
         dtype_constraint, Sequence
@@ -185,10 +201,12 @@ def _constrain_dtype(dataset, *, dtype_constraint=None):
             f"Expected dtype to be of type one of {dtype_constraint}, but got {dataset.dtype}."
         )
 
-    return dataset
-
 
+@_validator_from_condition
 def _constraint_dim(dataset, *, min_dim=None, max_dim=None):
+    if min_dim is not None and max_dim is not None and min_dim > max_dim:
+        raise ValueError("Impossible to satisfy dimension constraints on dataset.")
+
     min_dim = 0 if min_dim is None else min_dim
 
     dims = len(dataset.shape)
@@ -198,42 +216,26 @@ def _constraint_dim(dataset, *, min_dim=None, max_dim=None):
             f"Expected {min_dim} <= dimension of shape{f' <= {max_dim}'}, but got shape = {dataset.shape}."
         )
 
-    return dataset
-
 
+@_validator_from_condition
 def _constraint_shape(dataset, *, shape_constraint=None):
-    if shape_constraint and (
-        len(shape_constraint) != len(dataset.shape)
-        or reduce(
-            lambda x, y: x or y,
-            map(
-                lambda x: x[0] is not None and x[0] != x[1],
-                zip(shape_constraint, dataset.shape),
-            ),
-        )
-    ):
+    if shape_constraint and not _flex_shape_equal(shape_constraint, dataset.shape):
         raise ValueError(
             f"Expected shape to be {shape_constraint}, but got {dataset.shape}."
         )
 
-    return dataset
-
 
 def condataset(
     *, shape_constraint=None, dtype_constraint=None, min_dim=None, max_dim=None
 ):
     return Annotated[
-        Dataset,
-        BeforeValidator(Dataset.cast),
-        AfterValidator(partial(_constrain_dtype, dtype_constraint=dtype_constraint)),
-        AfterValidator(partial(_constraint_dim, min_dim=min_dim, max_dim=max_dim)),
-        AfterValidator(partial(_constraint_shape, shape_constraint=shape_constraint)),
+        CastDataset,
+        AfterValidator(_constrain_dtype(dtype_constraint=dtype_constraint)),
+        AfterValidator(_constraint_dim(min_dim=min_dim, max_dim=max_dim)),
+        AfterValidator(_constraint_shape(shape_constraint=shape_constraint)),
     ]
 
 
-CastDataset = Annotated[Dataset, BeforeValidator(Dataset.cast)]
-
-
 ########################################################################################
 
 
@@ -255,27 +257,46 @@ class GroupBase(BaseModel, extra="forbid"):
 
     attrs: Attrs = {}
 
+    @classmethod
+    def _is_allowed_field_type(cls, v):
+        is_dataset = Dataset._is_dataset_type(v)
+
+        is_annotated_dataset = typing.get_origin(
+            v
+        ) is Annotated and Dataset._is_dataset_type(v.__origin__)
+
+        is_optional_dataset = typing.get_origin(v) is Union and (
+            (v.__args__[0] == NoneType and Dataset._is_dataset_type(v.__args__[1]))
+            or (v.__args__[1] == NoneType and Dataset._is_dataset_type(v.__args__[0]))
+        )
+
+        is_dict_dataset = (
+            typing.get_origin(v) is dict
+            and v.__args__[0] is str
+            and Dataset._is_dataset_type(v.__args__[1])
+        )
+
+        return (
+            is_dataset or is_annotated_dataset or is_optional_dataset or is_dict_dataset
+        )
+
+    @classmethod
+    def _is_classvar(cls, v):
+        return v is ClassVar or typing.get_origin(v) is ClassVar
+
     def __init_subclass__(cls, **kwargs):
         super().__init_subclass__(**kwargs)
 
         for k, v in cls.__annotations__.items():
-            if k == "class_":
-                raise AttributeError("`class_` attribute should not be set manually.")
-
-            if k == "attrs" and k is not Attrs:
-                raise TypeError("`attrs` should be of type `Attrs`")
-
-            if (
-                k not in ["class_", "attrs"]
-                and v is not ClassVar
-                and not Dataset._is_dataset_type(v)
-                and not (typing.get_origin(v) is Annotated and v.__origin__ is Dataset)
-                and not (
-                    typing.get_origin(v) is dict
-                    and v.__args__[0] is str
-                    and Dataset._is_dataset_type(v.__args__[1])
+            if k in ["class_", "attrs"]:
+                raise AttributeError(
+                    "`class_` and `attrs` attribute should not be set manually."
                 )
-            ):
+
+            if cls._is_classvar(v):
+                continue
+
+            if not cls._is_allowed_field_type(v):
                 raise TypeError(
                     "All fields of `GroupBase` have to be of type `Dataset`."
                 )

src/oqd_dataschema/utils.py

@@ -0,0 +1,43 @@
+# Copyright 2024-2025 Open Quantum Design
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from functools import reduce
+
+########################################################################################
+
+__all__ = ["_flex_shape_equal", "_validator_from_condition"]
+
+
+########################################################################################
+
+
+def _flex_shape_equal(shape1, shape2):
+    return len(shape1) == len(shape2) and reduce(
+        lambda x, y: x and y,
+        map(
+            lambda x: x[0] is None or x[1] is None or x[0] == x[1],
+            zip(shape1, shape2),
+        ),
+    )
+
+
+def _validator_from_condition(f):
+    def _wrapped_validator(*args, **kwargs):
+        def _wrapped_condition(model):
+            f(model, *args, **kwargs)
+            return model
+
+        return _wrapped_condition
+
+    return _wrapped_validator