[feat] Implemented CastDataset that cast a numpy array automatically to a Dataset

Author: yhteoh

src/oqd_dataschema/__init__.py

@@ -33,4 +33,5 @@
     "OQDTestbenchDataGroup",
     "SinaraRawDataGroup",
     "condataset",
+    "CastDataset",
 ]

src/oqd_dataschema/base.py

@@ -33,7 +33,7 @@
 
 ########################################################################################
 
-__all__ = ["GroupBase", "Dataset", "GroupRegistry", "condataset"]
+__all__ = ["GroupBase", "Dataset", "GroupRegistry", "condataset", "CastDataset"]
 
 ########################################################################################
 
@@ -117,6 +117,12 @@ class Dataset(BaseModel, extra="forbid"):
         use_enum_values=False, arbitrary_types_allowed=True, validate_assignment=True
     )
 
+    @classmethod
+    def cast(cls, data):
+        if isinstance(data, np.ndarray):
+            return cls(data=data)
+        return data
+
     @model_validator(mode="before")
     @classmethod
     def validate_and_update(cls, values: dict):
@@ -212,12 +218,16 @@ def condataset(
 ):
     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 = Annotated[Dataset, BeforeValidator(Dataset.cast)]
+
+
 ########################################################################################
 
 

src/oqd_dataschema/groups.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 
-from oqd_dataschema.base import Dataset, GroupBase
+from oqd_dataschema.base import CastDataset, GroupBase
 
 ########################################################################################
 
@@ -33,7 +33,7 @@ class SinaraRawDataGroup(GroupBase):
     This is a placeholder for demonstration and development.
     """
 
-    camera_images: Dataset
+    camera_images: CastDataset
 
 
 class MeasurementOutcomesDataGroup(GroupBase):
@@ -42,7 +42,7 @@ class MeasurementOutcomesDataGroup(GroupBase):
     This is a placeholder for demonstration and development.
     """
 
-    outcomes: Dataset
+    outcomes: CastDataset
 
 
 class ExpectationValueDataGroup(GroupBase):
@@ -51,11 +51,11 @@ class ExpectationValueDataGroup(GroupBase):
     This is a placeholder for demonstration and development.
     """
 
-    expectation_value: Dataset
+    expectation_value: CastDataset
 
 
 class OQDTestbenchDataGroup(GroupBase):
     """ """
 
-    time: Dataset
-    voltages: Dataset
+    time: CastDataset
+    voltages: CastDataset