Merge pull request #7 from OpenQuantumDesign/refactor_group_registry

Refactor group registry

Author: yhteoh

.gitignore

@@ -172,4 +172,5 @@ cython_debug/
 
 *DS_Store
 *.h5
-*.code-workspace
+*.code-workspace
+.pre-commit-config.yaml

examples/custom_group.ipynb

@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 15,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -11,7 +11,7 @@
     "import numpy as np\n",
     "from rich.pretty import pprint\n",
     "\n",
-    "from oqd_dataschema.base import Dataset, Group, GroupRegistry\n",
+    "from oqd_dataschema.base import Dataset, GroupBase, GroupRegistry\n",
     "from oqd_dataschema.datastore import Datastore\n",
     "from oqd_dataschema.groups import (\n",
     "    SinaraRawDataGroup,\n",
@@ -20,20 +20,11 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 2,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "/Users/benjamin/Desktop/1 - Projects/Open Quantum Design/repos/oqd-dataschema/src/oqd_dataschema/base.py:66: UserWarning: Group type 'YourCustomGroup' is already registered. Overwriting <class '__main__.YourCustomGroup'> with <class '__main__.YourCustomGroup'>.\n",
-      "  GroupRegistry.register(cls)\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
-    "class YourCustomGroup(Group):\n",
+    "class YourCustomGroup(GroupBase):\n",
     "    \"\"\"\n",
     "    Here we define a custom Group, which is automatically added at runtime to the GroupRegistry.\n",
     "    \"\"\"\n",
@@ -43,31 +34,31 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 18,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "['SinaraRawDataGroup',\n",
-       " 'MeasurementOutcomesDataGroup',\n",
-       " 'ExpectationValueDataGroup',\n",
-       " 'OQDTestbenchDataGroup',\n",
-       " 'YourCustomGroup']"
+       "{'SinaraRawDataGroup': oqd_dataschema.groups.SinaraRawDataGroup,\n",
+       " 'MeasurementOutcomesDataGroup': oqd_dataschema.groups.MeasurementOutcomesDataGroup,\n",
+       " 'ExpectationValueDataGroup': oqd_dataschema.groups.ExpectationValueDataGroup,\n",
+       " 'OQDTestbenchDataGroup': oqd_dataschema.groups.OQDTestbenchDataGroup,\n",
+       " 'YourCustomGroup': __main__.YourCustomGroup}"
       ]
      },
-     "execution_count": 18,
+     "execution_count": 3,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
-    "GroupRegistry.list_types()"
+    "GroupRegistry.groups"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -82,7 +73,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": 5,
    "metadata": {},
    "outputs": [
     {
@@ -193,7 +184,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": 6,
    "metadata": {},
    "outputs": [
     {
@@ -321,7 +312,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.1"
+   "version": "3.13.2"
   }
  },
  "nbformat": 4,

examples/dataschema.ipynb

@@ -53,7 +53,10 @@ fixable = ["ALL"]
 
 [dependency-groups]
 dev = [
+    "jupyter>=1.1.1",
     "pre-commit>=4.1.0",
+    "rich>=14.1.0",
+    "ruff>=0.13.1",
 ]
 
 

pyproject.toml

@@ -0,0 +1,35 @@
+# 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 .base import Dataset, GroupBase, GroupRegistry
+from .datastore import Datastore
+from .groups import (
+    ExpectationValueDataGroup,
+    MeasurementOutcomesDataGroup,
+    OQDTestbenchDataGroup,
+    SinaraRawDataGroup,
+)
+
+########################################################################################
+
+__all__ = [
+    "Dataset",
+    "Datastore",
+    "GroupBase",
+    "GroupRegistry",
+    "ExpectationValueDataGroup",
+    "MeasurementOutcomesDataGroup",
+    "OQDTestbenchDataGroup",
+    "SinaraRawDataGroup",
+]

src/oqd_dataschema/__init__.py

@@ -13,18 +13,27 @@
 # limitations under the License.
 
 # %%
-from typing import Any, Literal, Optional, Type, Union
+import warnings
+from typing import Annotated, Any, Literal, Optional, Union
 
 import numpy as np
 from bidict import bidict
 from pydantic import (
     BaseModel,
     ConfigDict,
+    Discriminator,
     Field,
     TypeAdapter,
     model_validator,
 )
 
+########################################################################################
+
+__all__ = ["GroupBase", "Dataset", "GroupRegistry"]
+
+########################################################################################
+
+
 # %%
 mapping = bidict(
     {
@@ -39,7 +48,7 @@
 )
 
 
-class Group(BaseModel, extra="forbid"):
+class GroupBase(BaseModel, extra="forbid"):
     """
     Schema representation for a group object within an HDF5 file.
 
@@ -144,63 +153,40 @@ def validate_data_matches_shape_dtype(self):
         return self
 
 
-class GroupRegistry:
-    """Registry for managing group types dynamically"""
+class MetaGroupRegistry(type):
+    def __new__(cls, clsname, superclasses, attributedict):
+        attributedict["groups"] = dict()
+        return super().__new__(cls, clsname, superclasses, attributedict)
 
-    _types: dict[str, Type[Group]] = {}
-    _union_cache = None
+    def register(cls, group):
+        if not issubclass(group, GroupBase):
+            raise TypeError("You may only register subclasses of GroupBase.")
 
-    @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
+        if group.__name__ in cls.groups.keys():
+            warnings.warn(
+                f"Overwriting previously registered `{group.__name__}` group of the same name.",
+                UserWarning,
+                stacklevel=2,
+            )
 
-    @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")
+        cls.groups[group.__name__] = group
 
-            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)]
+    def clear(cls):
+        """Clear all registered types (useful for testing)"""
+        cls.groups.clear()
 
-        return cls._union_cache
+    @property
+    def union(cls):
+        """Get the current Union of all registered types"""
+        return Annotated[
+            Union[tuple(cls.groups.values())], Discriminator(discriminator="class_")
+        ]
 
-    @classmethod
-    def get_adapter(cls):
+    @property
+    def adapter(cls):
         """Get TypeAdapter for current registered types"""
-        from typing import Annotated
+        return TypeAdapter(cls.union)
 
-        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())
+class GroupRegistry(metaclass=MetaGroupRegistry):
+    pass

src/oqd_dataschema/base.py

@@ -22,7 +22,13 @@
 from pydantic import BaseModel, model_validator
 from pydantic.types import TypeVar
 
-from oqd_dataschema.base import Dataset, Group, GroupRegistry
+from oqd_dataschema.base import Dataset, GroupBase, GroupRegistry
+
+########################################################################################
+
+__all__ = ["Datastore"]
+
+########################################################################################
 
 
 # %%
@@ -44,16 +50,17 @@ 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):
+                    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] = adapter.validate_python(group_data)
+                        validated_groups[key] = GroupRegistry.adapter.validate_python(
+                            group_data
+                        )
                     else:
                         raise ValueError(
                             f"Invalid group data for key '{key}': {type(group_data)}"

src/oqd_dataschema/datastore.py

@@ -13,10 +13,21 @@
 # limitations under the License.
 
 
-from oqd_dataschema.base import Dataset, Group
+from oqd_dataschema.base import Dataset, GroupBase
 
+########################################################################################
 
-class SinaraRawDataGroup(Group):
+__all__ = [
+    "SinaraRawDataGroup",
+    "MeasurementOutcomesDataGroup",
+    "ExpectationValueDataGroup",
+    "OQDTestbenchDataGroup",
+]
+
+########################################################################################
+
+
+class SinaraRawDataGroup(GroupBase):
     """
     Example `Group` for raw data from the Sinara real-time control system.
     This is a placeholder for demonstration and development.
@@ -25,7 +36,7 @@ class SinaraRawDataGroup(Group):
     camera_images: Dataset
 
 
-class MeasurementOutcomesDataGroup(Group):
+class MeasurementOutcomesDataGroup(GroupBase):
     """
     Example `Group` for processed data classifying the readout of the state.
     This is a placeholder for demonstration and development.
@@ -34,7 +45,7 @@ class MeasurementOutcomesDataGroup(Group):
     outcomes: Dataset
 
 
-class ExpectationValueDataGroup(Group):
+class ExpectationValueDataGroup(GroupBase):
     """
     Example `Group` for processed data calculating the expectation values.
     This is a placeholder for demonstration and development.
@@ -43,7 +54,7 @@ class ExpectationValueDataGroup(Group):
     expectation_value: Dataset
 
 
-class OQDTestbenchDataGroup(Group):
+class OQDTestbenchDataGroup(GroupBase):
     """ """
 
     time: Dataset

src/oqd_dataschema/groups.py

@@ -18,7 +18,7 @@
 
 import numpy as np
 
-from oqd_dataschema.base import Dataset, Group
+from oqd_dataschema.base import Dataset, GroupBase
 from oqd_dataschema.datastore import Datastore
 from oqd_dataschema.groups import (
     SinaraRawDataGroup,
@@ -27,7 +27,7 @@
 
 # %%
 def test_adapt():
-    class TestNewGroup(Group):
+    class TestNewGroup(GroupBase):
         """ """
 
         array: Dataset