[fix, feat] Added method to Datastore to add and update groups. fixed datastore model_dump_hdf5 ignore attrs and class_ when iterating through fields

Author: yhteoh

src/oqd_dataschema/base.py

@@ -189,6 +189,7 @@ def _is_dataset_type(cls, type_):
 
 @_validator_from_condition
 def _constrain_dtype(dataset, *, dtype_constraint=None):
+    """Constrains the dtype of a dataset"""
     if (not isinstance(dtype_constraint, str)) and isinstance(
         dtype_constraint, Sequence
     ):
@@ -204,6 +205,7 @@ def _constrain_dtype(dataset, *, dtype_constraint=None):
 
 @_validator_from_condition
 def _constraint_dim(dataset, *, min_dim=None, max_dim=None):
+    """Constrains the dimension of a dataset"""
     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.")
 
@@ -219,6 +221,7 @@ def _constraint_dim(dataset, *, min_dim=None, max_dim=None):
 
 @_validator_from_condition
 def _constraint_shape(dataset, *, shape_constraint=None):
+    """Constrains the shape of a dataset"""
     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}."
@@ -228,6 +231,7 @@ def _constraint_shape(dataset, *, shape_constraint=None):
 def condataset(
     *, shape_constraint=None, dtype_constraint=None, min_dim=None, max_dim=None
 ):
+    """Implements dtype, dimension and shape constrains on the dataset."""
     return Annotated[
         CastDataset,
         AfterValidator(_constrain_dtype(dtype_constraint=dtype_constraint)),
@@ -312,11 +316,16 @@ def __init_subclass__(cls, **kwargs):
 
 
 class MetaGroupRegistry(type):
+    """
+    Metaclass for the GroupRegistry
+    """
+
     def __new__(cls, clsname, superclasses, attributedict):
         attributedict["groups"] = dict()
         return super().__new__(cls, clsname, superclasses, attributedict)
 
     def register(cls, group):
+        """Registers a group into the GroupRegistry."""
         if not issubclass(group, GroupBase):
             raise TypeError("You may only register subclasses of GroupBase.")
 
@@ -347,6 +356,10 @@ def adapter(cls):
 
 
 class GroupRegistry(metaclass=MetaGroupRegistry):
+    """
+    Represents the GroupRegistry
+    """
+
     pass
 
 

src/oqd_dataschema/datastore.py

@@ -37,12 +37,11 @@
 # %%
 class Datastore(BaseModel, extra="forbid"):
     """
-    Saves the model and its associated data to an HDF5 file.
-    This method serializes the model's data and attributes into an HDF5 file
-    at the specified filepath.
+    Class representing a datastore with restricted HDF5 format.
 
     Attributes:
-        filepath (pathlib.Path): The path to the HDF5 file where the model data will be saved.
+        groups (Dict[str,Group]): groups of data.
+        attrs (Attrs): attributes of the datastore.
     """
 
     groups: Dict[str, Any]
@@ -51,6 +50,7 @@ class Datastore(BaseModel, extra="forbid"):
 
     @classmethod
     def _validate_group(cls, key, group):
+        """Helper function for validating group to be of type Group registered in the GroupRegistry."""
         if isinstance(group, GroupBase):
             return group
 
@@ -62,6 +62,7 @@ def _validate_group(cls, key, group):
     @field_validator("groups", mode="before")
     @classmethod
     def validate_groups(cls, data):
+        """Validates groups to be of type Group registered in the GroupRegistry."""
         if GroupRegistry.groups == {}:
             raise ValueError(
                 "No group types available. Register group types before creating Datastore."
@@ -71,6 +72,7 @@ def validate_groups(cls, data):
         return validated_groups
 
     def _dump_group(self, h5datastore, gkey, group):
+        """Helper function for dumping Group."""
         # remove existing group
         if gkey in h5datastore.keys():
             del h5datastore[gkey]
@@ -89,15 +91,20 @@ def _dump_group(self, h5datastore, gkey, group):
 
         # dump group data
         for dkey, dataset in group.__dict__.items():
+            if dkey in ["attr", "class_"]:
+                continue
+
             # if group field contain dictionary of Dataset
             if isinstance(dataset, dict):
                 h5_subgroup = h5_group.create_group(dkey)
                 for ddkey, ddataset in dataset.items():
                     self._dump_dataset(h5_subgroup, ddkey, ddataset)
-            else:
-                self._dump_dataset(h5_group, dkey, dataset)
+                continue
+
+            self._dump_dataset(h5_group, dkey, dataset)
 
     def _dump_dataset(self, h5group, dkey, dataset):
+        """Helper function for dumping Dataset."""
         if not isinstance(dataset, Dataset):
             raise ValueError("Group data field is not a Dataset.")
 
@@ -132,6 +139,9 @@ def model_dump_hdf5(self, filepath: pathlib.Path, mode: Literal["w", "a"] = "w")
 
             # dump each group
             for gkey, group in self.groups.items():
+                if gkey in ["attr", "class_"]:
+                    continue
+
                 self._dump_group(f, gkey, group)
 
     @classmethod
@@ -177,7 +187,26 @@ def model_validate_hdf5(cls, filepath: pathlib.Path):
             return self
 
     def __getitem__(self, key):
+        """Overloads indexing to retrieve elements in groups."""
         return self.groups.__getitem__(key)
 
     def __iter__(self):
+        """Overloads iter to iterate over elements in groups."""
         return self.groups.items().__iter__()
+
+    def add(self, **groups):
+        """Adds a new groups to the datastore."""
+        for k, v in groups.items():
+            if k in self.groups.keys():
+                raise ValueError(
+                    "Key already exist in the datastore, use `update` instead if intending to overwrite past data."
+                )
+            self.groups[k] = v
+
+    def update(self, **groups):
+        """Updates groups in the datastore, overwriting past values."""
+        for k, v in groups.items():
+            self.groups[k] = v
+
+
+# %%

src/oqd_dataschema/utils.py

@@ -23,6 +23,7 @@
 
 
 def _flex_shape_equal(shape1, shape2):
+    """Helper function for comparing concrete and flex shapes."""
     return len(shape1) == len(shape2) and reduce(
         lambda x, y: x and y,
         map(
@@ -33,6 +34,8 @@ def _flex_shape_equal(shape1, shape2):
 
 
 def _validator_from_condition(f):
+    """Helper decorator for turning a condition into a validation."""
+
     def _wrapped_validator(*args, **kwargs):
         def _wrapped_condition(model):
             f(model, *args, **kwargs)