Wip experiment with model for metadata

Author: jenshnielsen

pyproject.toml

@@ -56,6 +56,7 @@ dependencies = [
     "tornado>=6.3.3",
     "ipython>=8.10.0",
     "pillow>=9.0.0",
+    "pydantic>=2.0.0",
 ]
 
 dynamic = ["version"]

src/qcodes/metadatable/metadatable_base.py

@@ -1,5 +1,8 @@
 from abc import abstractmethod
-from typing import TYPE_CHECKING, Any, Optional, final
+from typing import TYPE_CHECKING, Any, Generic, Optional, final
+
+from pydantic import BaseModel
+from typing_extensions import TypeVar
 
 from qcodes.utils import deep_update
 
@@ -18,9 +21,29 @@
 Snapshot = dict[str, Any]
 
 
-class Metadatable:
-    def __init__(self, metadata: Optional["Mapping[str, Any]"] = None):
+class EmptyMetaDataModel(BaseModel):
+    pass
+
+MetaDataSnapShotType = TypeVar("MetaDataSnapShotType", bound=EmptyMetaDataModel)
+
+
+class EmptyTypedSnapShot(BaseModel):
+    pass
+
+
+SnapShotType = TypeVar("SnapShotType", bound=EmptyTypedSnapShot)
+
+
+class Metadatable(Generic[SnapShotType, MetaDataSnapShotType]):
+    def __init__(
+        self,
+        metadata: Optional["Mapping[str, Any]"] = None,
+        model: type[SnapShotType] = EmptyTypedSnapShot,
+        metadata_model: type[MetaDataSnapShotType] = EmptyMetaDataModel,
+    ):
         self.metadata: dict[str, Any] = {}
+        self._model = model or EmptyTypedSnapShot
+        self._metadata_model = metadata_model or EmptyMetaDataModel
         self.load_metadata(metadata or {})
 
     def load_metadata(self, metadata: "Mapping[str, Any]") -> None:
@@ -53,6 +76,16 @@ def snapshot(self, update: Optional[bool] = False) -> Snapshot:
 
         return snap
 
+    @final
+    def typed_snapshot(self) -> SnapShotType:
+        snapshot_dict = self.snapshot()  # probably want to filter metadata here
+        snapshot = self._model(**snapshot_dict)
+        return snapshot
+
+    @final
+    def typed_metadata(self) -> MetaDataSnapShotType:
+        return self._metadata_model(**self.metadata)
+
     def snapshot_base(
         self,
         update: Optional[bool] = False,
@@ -63,8 +96,20 @@ def snapshot_base(
         """
         return {}
 
+    # @property
+    # def metadata_model(self) -> type[BaseModel] | None:
+    #     return self._metadata_model
+
+
+    # @metadata_model.setter
+    # def metadata_model(self, model: type[BaseModel] | None) -> None:
+    #     self._metadata_model = model
+
 
-class MetadatableWithName(Metadatable):
+class MetadatableWithName(
+    Metadatable[SnapShotType, MetaDataSnapShotType],
+    Generic[SnapShotType, MetaDataSnapShotType],
+):
     """Add short_name and full_name properties to Metadatable.
     This is used as a base class for all components in QCoDeS that
     are members of a station to ensure that they have a name and

src/qcodes/parameters/parameter.py

@@ -6,7 +6,14 @@
 import logging
 import os
 from types import MethodType
-from typing import TYPE_CHECKING, Any, Callable, Literal
+from typing import TYPE_CHECKING, Any, Callable, Generic, Literal
+
+from qcodes.metadatable.metadatable_base import (
+    EmptyMetaDataModel,
+    EmptyTypedSnapShot,
+    MetaDataSnapShotType,
+    SnapShotType,
+)
 
 from .command import Command
 from .parameter_base import ParamDataType, ParameterBase, ParamRawDataType
@@ -21,7 +28,10 @@
 log = logging.getLogger(__name__)
 
 
-class Parameter(ParameterBase):
+class Parameter(
+    ParameterBase[SnapShotType, MetaDataSnapShotType],
+    Generic[SnapShotType, MetaDataSnapShotType],
+):
     """
     A parameter represents a single degree of freedom. Most often,
     this is the standard parameter for Instruments, though it can also be
@@ -179,6 +189,8 @@ def __init__(
         docstring: str | None = None,
         initial_cache_value: float | str | None = None,
         bind_to_instrument: bool = True,
+        model: type[SnapShotType] = EmptyTypedSnapShot,
+        metadata_model: type[MetaDataSnapShotType] = EmptyMetaDataModel,
         **kwargs: Any,
     ) -> None:
         def _get_manual_parameter(self: Parameter) -> ParamRawDataType:
@@ -240,6 +252,8 @@ def _set_manual_parameter(
             vals=vals,
             max_val_age=max_val_age,
             bind_to_instrument=bind_to_instrument,
+            model=model,
+            metadata_model=metadata_model,
             **kwargs,
         )
 
@@ -441,6 +455,19 @@ def sweep(
         return SweepFixedValues(self, start=start, stop=stop, step=step, num=num)
 
 
+class ParameterSnapshot(EmptyTypedSnapShot):
+    # need to handle replacing __class__ with a different name that is compatible
+    value: Any  # use paramdatatype
+    raw_value: Any  # useparamdatatype
+    ts: str | None
+    inter_delay: float
+    name: str
+    post_delay: float
+    validators: list[str]
+    label: str
+    unit: str
+
+
 class ManualParameter(Parameter):
     def __init__(
         self,

src/qcodes/parameters/parameter_base.py

@@ -7,9 +7,18 @@
 from contextlib import contextmanager
 from datetime import datetime
 from functools import cached_property, wraps
-from typing import TYPE_CHECKING, Any, ClassVar, overload
-
-from qcodes.metadatable import Metadatable, MetadatableWithName
+from typing import TYPE_CHECKING, Any, ClassVar, Generic, overload
+
+from qcodes.metadatable import (
+    Metadatable,
+    MetadatableWithName,
+)
+from qcodes.metadatable.metadatable_base import (
+    EmptyMetaDataModel,
+    EmptyTypedSnapShot,
+    MetaDataSnapShotType,
+    SnapShotType,
+)
 from qcodes.utils import DelegateAttributes, full_class, qcodes_abstractmethod
 from qcodes.validators import Enum, Ints, Validator
 
@@ -85,7 +94,10 @@ def invert_val_mapping(val_mapping: Mapping[Any, Any]) -> dict[Any, Any]:
     return {v: k for k, v in val_mapping.items()}
 
 
-class ParameterBase(MetadatableWithName):
+class ParameterBase(
+    MetadatableWithName[SnapShotType, MetaDataSnapShotType],
+    Generic[SnapShotType, MetaDataSnapShotType],
+):
     """
     Shared behavior for all parameters. Not intended to be used
     directly, normally you should use ``Parameter``, ``ArrayParameter``,
@@ -201,8 +213,11 @@ def __init__(
         abstract: bool | None = False,
         bind_to_instrument: bool = True,
         register_name: str | None = None,
+        model: type[SnapShotType] = EmptyTypedSnapShot,
+        metadata_model: type[MetaDataSnapShotType] = EmptyMetaDataModel,
+        **kwargs: Any,
     ) -> None:
-        super().__init__(metadata)
+        super().__init__(metadata, model=model, metadata_model=metadata_model, **kwargs)
         if not str(name).isidentifier():
             raise ValueError(
                 f"Parameter name must be a valid identifier "