Merge pull request #5 from OpenQuantumDesign/registry

Add registry for dynamic type definition in Datastore

Author: benjimaclellan

examples/custom_group.ipynb

@@ -501,7 +501,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": ".venv",
+   "display_name": "oqd-dataschema",
    "language": "python",
    "name": "python3"
   },

examples/test_h5pydantic.ipynb renamed to examples/dataschema.ipynb

@@ -13,14 +13,15 @@
 # limitations under the License.
 
 # %%
-from typing import Any, Literal, Optional, Union
+from typing import Any, Literal, Optional, Type, Union
 
 import numpy as np
 from bidict import bidict
 from pydantic import (
     BaseModel,
     ConfigDict,
     Field,
+    TypeAdapter,
     model_validator,
 )
 
@@ -38,7 +39,7 @@
 )
 
 
-class Group(BaseModel):
+class Group(BaseModel, extra="forbid"):
     """
     Schema representation for a group object within an HDF5 file.
 
@@ -56,8 +57,25 @@ class Group(BaseModel):
 
     attrs: Optional[dict[str, Union[int, float, str, complex]]] = {}
 
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+        cls.__annotations__["class_"] = Literal[cls.__name__]
+        setattr(cls, "class_", cls.__name__)
 
-class Dataset(BaseModel):
+        # Auto-register new group types
+        GroupRegistry.register(cls)
+
+    @model_validator(mode="before")
+    @classmethod
+    def auto_assign_class(cls, data):
+        if isinstance(data, BaseModel):
+            return data
+        if isinstance(data, dict):
+            data["class_"] = cls.__name__
+        return data
+
+
+class Dataset(BaseModel, extra="forbid"):
     """
     Schema representation for a dataset object to be saved within an HDF5 file.
 
@@ -110,12 +128,6 @@ def validate_and_update(cls, values: dict):
 
         return values
 
-        # else:
-        #     assert data.dtype == dtype and data.shape == shape
-
-        # else:
-        #     raise ValueError("Must provide either `dtype` and `shape` or `data`.")
-
     @model_validator(mode="after")
     def validate_data_matches_shape_dtype(self):
         """Ensure that `data` matches `dtype` and `shape`."""
@@ -130,3 +142,65 @@ def validate_data_matches_shape_dtype(self):
                     f"Expected shape {self.shape}, but got {self.data.shape}."
                 )
         return self
+
+
+class GroupRegistry:
+    """Registry for managing group types dynamically"""
+
+    _types: dict[str, Type[Group]] = {}
+    _union_cache = None
+
+    @classmethod
+    def register(cls, group_type: Type[Group]):
+        """Register a new group type"""
+        import warnings
+
+        type_name = group_type.__name__
+
+        # Check if type is already registered
+        if type_name in cls._types:
+            existing_type = cls._types[type_name]
+            if existing_type is not group_type:  # Different class with same name
+                warnings.warn(
+                    f"Group type '{type_name}' is already registered. "
+                    f"Overwriting {existing_type} with {group_type}.",
+                    UserWarning,
+                    stacklevel=2,
+                )
+
+        cls._types[type_name] = group_type
+        cls._union_cache = None  # Invalidate cache
+
+    @classmethod
+    def get_union(cls):
+        """Get the current Union of all registered types"""
+        if cls._union_cache is None:
+            if not cls._types:
+                raise ValueError("No group types registered")
+
+            type_list = list(cls._types.values())
+            if len(type_list) == 1:
+                cls._union_cache = type_list[0]
+            else:
+                cls._union_cache = Union[tuple(type_list)]
+
+        return cls._union_cache
+
+    @classmethod
+    def get_adapter(cls):
+        """Get TypeAdapter for current registered types"""
+        from typing import Annotated
+
+        union_type = cls.get_union()
+        return TypeAdapter(Annotated[union_type, Field(discriminator="class_")])
+
+    @classmethod
+    def clear(cls):
+        """Clear all registered types (useful for testing)"""
+        cls._types.clear()
+        cls._union_cache = None
+
+    @classmethod
+    def list_types(cls):
+        """List all registered type names"""
+        return list(cls._types.keys())

src/oqd_dataschema/base.py

@@ -12,26 +12,21 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# %%
+
 import pathlib
-from typing import Union
+from typing import Any, Dict, Literal, Optional
 
 import h5py
 import numpy as np
-from pydantic import BaseModel
-
-from oqd_dataschema.base import Dataset
-from oqd_dataschema.groups import (
-    ExpectationValueDataGroup,
-    MeasurementOutcomesDataGroup,
-    SinaraRawDataGroup,
-)
+from pydantic import BaseModel, model_validator
+from pydantic.types import TypeVar
 
-GroupSubtypes = Union[
-    SinaraRawDataGroup, MeasurementOutcomesDataGroup, ExpectationValueDataGroup
-]
+from oqd_dataschema.base import Dataset, Group, GroupRegistry
 
 
-class Datastore(BaseModel):
+# %%
+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
@@ -41,9 +36,41 @@ class Datastore(BaseModel):
         filepath (pathlib.Path): The path to the HDF5 file where the model data will be saved.
     """
 
-    groups: dict[str, GroupSubtypes]
+    groups: Dict[str, Any]
 
-    def model_dump_hdf5(self, filepath: pathlib.Path):
+    @model_validator(mode="before")
+    @classmethod
+    def validate_groups(cls, data):
+        if isinstance(data, dict) and "groups" in data:
+            # Get the current adapter from registry
+            try:
+                adapter = GroupRegistry.get_adapter()
+                validated_groups = {}
+
+                for key, group_data in data["groups"].items():
+                    if isinstance(group_data, Group):
+                        # Already a Group instance
+                        validated_groups[key] = group_data
+                    elif isinstance(group_data, dict):
+                        # Parse dict using discriminated union
+                        validated_groups[key] = adapter.validate_python(group_data)
+                    else:
+                        raise ValueError(
+                            f"Invalid group data for key '{key}': {type(group_data)}"
+                        )
+
+                data["groups"] = validated_groups
+
+            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
+
+        return data
+
+    def model_dump_hdf5(self, filepath: pathlib.Path, mode: Literal["w", "a"] = "a"):
         """
         Saves the model and its associated data to an HDF5 file.
         This method serializes the model's data and attributes into an HDF5 file
@@ -54,8 +81,11 @@ def model_dump_hdf5(self, filepath: pathlib.Path):
         """
         filepath.parent.mkdir(exist_ok=True, parents=True)
 
-        with h5py.File(filepath, "a") as f:
+        with h5py.File(filepath, mode) as f:
+            # store the model JSON schema
             f.attrs["model"] = self.model_dump_json()
+
+            # store each group
             for gkey, group in self.groups.items():
                 if gkey in f.keys():
                     del f[gkey]
@@ -71,7 +101,9 @@ def model_dump_hdf5(self, filepath: pathlib.Path):
                         h5_dataset.attrs[akey] = attr
 
     @classmethod
-    def model_validate_hdf5(cls, filepath: pathlib.Path):
+    def model_validate_hdf5(
+        cls, filepath: pathlib.Path, types: Optional[TypeVar] = None
+    ):
         """
         Loads the model from an HDF5 file at the specified filepath.
 
@@ -80,9 +112,11 @@ def model_validate_hdf5(cls, filepath: pathlib.Path):
         """
         with h5py.File(filepath, "r") as f:
             self = cls.model_validate_json(f.attrs["model"])
+
+            # loop through all groups in the model schema and load HDF5 store
             for gkey, group in self.groups.items():
                 for dkey, val in group.__dict__.items():
-                    if dkey == "attrs":
+                    if dkey in ("attrs", "class_"):
                         continue
                     group.__dict__[dkey].data = np.array(f[gkey][dkey][()])
             return self

src/oqd_dataschema/datastore.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+
 from oqd_dataschema.base import Dataset, Group
 
 
@@ -43,8 +44,7 @@ class ExpectationValueDataGroup(Group):
 
 
 class OQDTestbenchDataGroup(Group):
-    """
-    """
+    """ """
 
     time: Dataset
-    voltages: Dataset
+    voltages: Dataset

src/oqd_dataschema/groups.py

@@ -48,3 +48,6 @@ def test_serialize_deserialize(dtype):
     data_reload = Datastore.model_validate_hdf5(filepath)
 
     assert data_reload.groups["test"].camera_images.data.dtype == mapping[dtype]
+
+
+# %%

tests/test_datastore.py

@@ -0,0 +1,51 @@
+# 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.
+
+
+# %%
+import pathlib
+
+import numpy as np
+
+from oqd_dataschema.base import Dataset, Group
+from oqd_dataschema.datastore import Datastore
+from oqd_dataschema.groups import (
+    SinaraRawDataGroup,
+)
+
+
+# %%
+def test_adapt():
+    class TestNewGroup(Group):
+        """ """
+
+        array: Dataset
+
+    filepath = pathlib.Path("test.h5")
+
+    data = np.ones([10, 10]).astype("int64")
+    group1 = TestNewGroup(array=Dataset(data=data))
+
+    data = np.ones([10, 10]).astype("int32")
+    group2 = SinaraRawDataGroup(camera_images=Dataset(data=data))
+
+    datastore = Datastore(
+        groups={
+            "group1": group1,
+            "group2": group2,
+        }
+    )
+    datastore.model_dump_hdf5(filepath, mode="w")
+
+    Datastore.model_validate_hdf5(filepath)