[WIP] Adding model to metadata and snapshot

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
+
+MetadataType = TypeVar("MetadataType", bound=EmptyMetadataModel)
+
+
+class EmptySnapshotModel(BaseModel):
+    pass
+
+
+SnapshotType = TypeVar("SnapshotType", bound=EmptySnapshotModel)
+
+
+class Metadatable(Generic[SnapshotType, MetadataType]):
+    def __init__(
+        self,
+        metadata: Optional["Mapping[str, Any]"] = None,
+        snapshot_model: type[SnapshotType] = EmptySnapshotModel,
+        metadata_model: type[MetadataType] = EmptyMetadataModel,
+    ):
         self.metadata: dict[str, Any] = {}
+        self._snapshot_model = snapshot_model or EmptySnapshotModel
+        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._snapshot_model(**snapshot_dict)
+        return snapshot
+
+    @final
+    def typed_metadata(self) -> MetadataType:
+        return self._metadata_model(**self.metadata)
+
     def snapshot_base(
         self,
         update: Optional[bool] = False,
@@ -63,8 +96,17 @@ 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, MetadataType]):
     """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

@@ -8,6 +8,13 @@
 from types import MethodType
 from typing import TYPE_CHECKING, Any, Callable, Literal
 
+from qcodes.metadatable.metadatable_base import (
+    EmptyMetadataModel,
+    EmptySnapshotModel,
+    MetadataType,
+    SnapshotType,
+)
+
 from .command import Command
 from .parameter_base import ParamDataType, ParameterBase, ParamRawDataType
 from .sweep_values import SweepFixedValues
@@ -21,7 +28,7 @@
 log = logging.getLogger(__name__)
 
 
-class Parameter(ParameterBase):
+class Parameter(ParameterBase[SnapshotType, MetadataType]):
     """
     A parameter represents a single degree of freedom. Most often,
     this is the standard parameter for Instruments, though it can also be
@@ -179,6 +186,8 @@ def __init__(
         docstring: str | None = None,
         initial_cache_value: float | str | None = None,
         bind_to_instrument: bool = True,
+        snapshot_model: type[SnapshotType] = EmptySnapshotModel,
+        metadata_model: type[MetadataType] = EmptyMetadataModel,
         **kwargs: Any,
     ) -> None:
         def _get_manual_parameter(self: Parameter) -> ParamRawDataType:
@@ -240,6 +249,8 @@ def _set_manual_parameter(
             vals=vals,
             max_val_age=max_val_age,
             bind_to_instrument=bind_to_instrument,
+            snapshot_model=snapshot_model,
+            metadata_model=metadata_model,
             **kwargs,
         )
 
@@ -441,6 +452,19 @@ def sweep(
         return SweepFixedValues(self, start=start, stop=stop, step=step, num=num)
 
 
+class ParameterSnapshot(EmptySnapshotModel):
+    # 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

@@ -9,7 +9,16 @@
 from functools import cached_property, wraps
 from typing import TYPE_CHECKING, Any, ClassVar, overload
 
-from qcodes.metadatable import Metadatable, MetadatableWithName
+from qcodes.metadatable import (
+    Metadatable,
+    MetadatableWithName,
+)
+from qcodes.metadatable.metadatable_base import (
+    EmptyMetadataModel,
+    EmptySnapshotModel,
+    MetadataType,
+    SnapshotType,
+)
 from qcodes.utils import DelegateAttributes, full_class, qcodes_abstractmethod
 from qcodes.validators import Enum, Ints, Validator
 
@@ -85,7 +94,7 @@ 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, MetadataType]):
     """
     Shared behavior for all parameters. Not intended to be used
     directly, normally you should use ``Parameter``, ``ArrayParameter``,
@@ -201,8 +210,11 @@ def __init__(
         abstract: bool | None = False,
         bind_to_instrument: bool = True,
         register_name: str | None = None,
+        snapshot_model: type[SnapshotType] = EmptySnapshotModel,
+        metadata_model: type[MetadataType] = EmptyMetadataModel,
+        **kwargs: Any,
     ) -> None:
-        super().__init__(metadata)
+        super().__init__(metadata, snapshot_model=snapshot_model, metadata_model=metadata_model, **kwargs)
         if not str(name).isidentifier():
             raise ValueError(
                 f"Parameter name must be a valid identifier "

tests/parameter/test_parameter_typed_metadata.py

@@ -0,0 +1,43 @@
+from enum import StrEnum
+
+from typing_extensions import assert_type
+
+from qcodes.metadatable.metadatable_base import EmptyMetaDataModel
+from qcodes.parameters import Parameter
+from qcodes.parameters.parameter import ParameterSnapshot
+
+
+def test_parameter_typed_metadata_basic():
+
+    class ParamType(StrEnum):
+
+        voltage = "voltage"
+        current = "current"
+
+    class MyParameterMetadata(EmptyMetaDataModel):
+        param_type: ParamType
+
+    a = Parameter(
+        name="myparam",
+        set_cmd=None,
+        get_cmd=None,
+        model=ParameterSnapshot,
+        metadata_model=MyParameterMetadata,
+    )
+    a.metadata["param_type"] = (
+        ParamType.voltage
+    )  # TODO setting metadata should validate against the model
+    value = 123
+    a.set(value)
+
+    assert isinstance(a.typed_snapshot(), ParameterSnapshot)
+    assert isinstance(a.typed_metadata(), MyParameterMetadata)
+
+    assert a.typed_metadata().param_type == ParamType.voltage
+    assert a.typed_snapshot().value == value
+    assert a.typed_snapshot().name == "myparam"
+
+    assert_type(
+        a.typed_metadata(), MyParameterMetadata
+    )  # TODO this is only checked if the type checker runs agains the test
+    assert_type(a.typed_snapshot(), ParameterSnapshot)