[refactor] refactor datastore group validation and added comments

Author: yhteoh

src/oqd_dataschema/datastore.py

@@ -49,69 +49,69 @@ class Datastore(BaseModel, extra="forbid"):
 
     attrs: Attrs = {}
 
-    @field_validator("groups", mode="before")
     @classmethod
-    def validate_groups(cls, data):
-        if isinstance(data, dict):
-            # Get the current adapter from registry
-            try:
-                validated_groups = {}
-
-                for key, group_data in data.items():
-                    if isinstance(group_data, GroupBase):
-                        # Already a Group instance
-                        validated_groups[key] = group_data
-                    elif isinstance(group_data, dict):
-                        # Parse dict using discriminated union
-                        validated_groups[key] = GroupRegistry.adapter.validate_python(
-                            group_data
-                        )
-                    else:
-                        raise ValueError(
-                            f"Invalid group data for key '{key}': {type(group_data)}"
-                        )
+    def _validate_group(cls, key, group):
+        if isinstance(group, GroupBase):
+            return group
 
-                data = validated_groups
+        if isinstance(group, dict):
+            return GroupRegistry.adapter.validate_python(group)
 
-            except ValueError as e:
-                if "No group types registered" in str(e):
-                    raise ValueError(
-                        "No group types available. Register group types before creating Datastore."
-                    )
-                raise
+        raise ValueError(f"Key `{key}` contains invalid group data.")
 
-        return data
+    @field_validator("groups", mode="before")
+    @classmethod
+    def validate_groups(cls, data):
+        if GroupRegistry.groups == {}:
+            raise ValueError(
+                "No group types available. Register group types before creating Datastore."
+            )
+
+        validated_groups = {k: cls._validate_group(k, v) for k, v in data.items()}
+        return validated_groups
 
     def _dump_group(self, h5datastore, gkey, group):
+        # remove existing group
         if gkey in h5datastore.keys():
             del h5datastore[gkey]
+
+        # create group
         h5_group = h5datastore.create_group(gkey)
 
+        # dump group schema
         h5_group.attrs["_group_schema"] = json.dumps(
             group.model_json_schema(), indent=2
         )
+
+        # dump group attributes
         for akey, attr in group.attrs.items():
             h5_group.attrs[akey] = attr
 
+        # dump group data
         for dkey, dataset in group.__dict__.items():
+            # 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)
-
-            self._dump_dataset(h5_group, dkey, dataset)
+            else:
+                self._dump_dataset(h5_group, dkey, dataset)
 
     def _dump_dataset(self, h5group, dkey, dataset):
-        if isinstance(dataset, Dataset):
-            if dataset.dtype in "str":
-                h5_dataset = h5group.create_dataset(
-                    dkey, data=dataset.data.astype(np.dtypes.BytesDType)
-                )
-            else:
-                h5_dataset = h5group.create_dataset(dkey, data=dataset.data)
+        if not isinstance(dataset, Dataset):
+            raise ValueError("Group data field is not a Dataset.")
+
+        # dtype str converted to bytes when dumped (h5 compatibility)
+        if dataset.dtype in "str":
+            h5_dataset = h5group.create_dataset(
+                dkey, data=dataset.data.astype(np.dtypes.BytesDType)
+            )
+        else:
+            h5_dataset = h5group.create_dataset(dkey, data=dataset.data)
 
-            for akey, attr in dataset.attrs.items():
-                h5_dataset.attrs[akey] = attr
+        # dump dataset attributes
+        for akey, attr in dataset.attrs.items():
+            h5_dataset.attrs[akey] = attr
 
     def model_dump_hdf5(self, filepath: pathlib.Path, mode: Literal["w", "a"] = "w"):
         """
@@ -125,12 +125,12 @@ def model_dump_hdf5(self, filepath: pathlib.Path, mode: Literal["w", "a"] = "w")
         filepath.parent.mkdir(exist_ok=True, parents=True)
 
         with h5py.File(filepath, mode) as f:
-            # store the model JSON schema
+            # dump the datastore signature
             f.attrs["_datastore_signature"] = self.model_dump_json(indent=2)
             for akey, attr in self.attrs.items():
                 f.attrs[akey] = attr
 
-            # store each group
+            # dump each group
             for gkey, group in self.groups.items():
                 self._dump_group(f, gkey, group)
 
@@ -143,26 +143,36 @@ def model_validate_hdf5(cls, filepath: pathlib.Path):
             filepath (pathlib.Path): The path to the HDF5 file where the model data will be read and validated from.
         """
         with h5py.File(filepath, "r") as f:
+            # Load datastore signature
             self = cls.model_validate_json(f.attrs["_datastore_signature"])
 
-            # loop through all groups in the model schema and load HDF5 store
+            # loop through all groups in the model schema and load the data
             for gkey, group in self:
                 for dkey in group.__class__.model_fields:
+                    # ignore attrs and class_ fields
                     if dkey in ("attrs", "class_"):
                         continue
 
+                    # load Dataset data
                     if isinstance(group.__dict__[dkey], Dataset):
                         group.__dict__[dkey].data = np.array(f[gkey][dkey][()]).astype(
                             DTypes.get(group.__dict__[dkey].dtype).value
                         )
+                        continue
 
+                    # load data for dict of Dataset
                     if isinstance(group.__dict__[dkey], dict):
                         for ddkey in group.__dict__[dkey]:
                             group.__dict__[dkey][ddkey].data = np.array(
                                 f[gkey][dkey][ddkey][()]
                             ).astype(
                                 DTypes.get(group.__dict__[dkey][ddkey].dtype).value
                             )
+                        continue
+
+                    raise TypeError(
+                        "Group data fields must be of type Dataset or dict of Dataset."
+                    )
 
             return self