Check if group in registry, union_cache, and forbid extra

Author: benjimaclellan

examples/custom_group.ipynb

@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 15,
    "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\n",
+    "from oqd_dataschema.base import Dataset, Group, GroupRegistry\n",
     "from oqd_dataschema.datastore import Datastore\n",
     "from oqd_dataschema.groups import (\n",
     "    SinaraRawDataGroup,\n",
@@ -20,9 +20,18 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "metadata": {},
-   "outputs": [],
+   "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"
+     ]
+    }
+   ],
    "source": [
     "class YourCustomGroup(Group):\n",
     "    \"\"\"\n",
@@ -34,7 +43,31 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 18,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "['SinaraRawDataGroup',\n",
+       " 'MeasurementOutcomesDataGroup',\n",
+       " 'ExpectationValueDataGroup',\n",
+       " 'OQDTestbenchDataGroup',\n",
+       " 'YourCustomGroup']"
+      ]
+     },
+     "execution_count": 18,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "GroupRegistry.list_types()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -49,7 +82,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 13,
    "metadata": {},
    "outputs": [
     {
@@ -160,7 +193,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 14,
    "metadata": {},
    "outputs": [
     {
@@ -263,6 +296,13 @@
     "parse = Datastore.model_validate_hdf5(filepath)\n",
     "pprint(parse)"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {

src/oqd_dataschema/base.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 # %%
-from typing import Any, Literal, Optional, Type, Union
+from typing import Any, Dict, Literal, Optional, Type, Union
 
 import numpy as np
 from bidict import bidict
@@ -39,7 +39,7 @@
 )
 
 
-class Group(BaseModel):
+class Group(BaseModel, extra="forbid"):
     """
     Schema representation for a group object within an HDF5 file.
 
@@ -75,7 +75,7 @@ def auto_assign_class(cls, data):
         return data
 
 
-class Dataset(BaseModel):
+class Dataset(BaseModel, extra="forbid"):
     """
     Schema representation for a dataset object to be saved within an HDF5 file.
 
@@ -147,29 +147,45 @@ def validate_data_matches_shape_dtype(self):
 class GroupRegistry:
     """Registry for managing group types dynamically"""
 
-    _types: dict[str, Type[Group]] = {}
-    _union_cache = None
+    _types: Dict[str, Type[Group]] = {}
 
     @classmethod
     def register(cls, group_type: Type[Group]):
         """Register a new group type"""
-        cls._types[group_type.__name__] = group_type
-        cls._union_cache = None  # Invalidate cache
+        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
+
+    @classmethod
+    @property
+    def union_cache(cls):
+        """Get the current Union of all registered types (computed on demand)"""
+        if not cls._types:
+            raise ValueError("No group types registered")
+
+        type_list = list(cls._types.values())
+        if len(type_list) == 1:
+            return type_list[0]
+        else:
+            return Union[tuple(type_list)]
 
     @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
+        return cls.union_cache
 
     @classmethod
     def get_adapter(cls):
@@ -183,7 +199,6 @@ def get_adapter(cls):
     def clear(cls):
         """Clear all registered types (useful for testing)"""
         cls._types.clear()
-        cls._union_cache = None
 
     @classmethod
     def list_types(cls):

src/oqd_dataschema/datastore.py

@@ -26,7 +26,7 @@
 
 
 # %%
-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