# Repository: https://gitlab.com/quantify-os/quantify-scheduler
# Licensed according to the LICENCE file on the main branch
"""Module containing quantify YAML utilities."""
from __future__ import annotations
import os
from datetime import datetime, timezone
from enum import Enum
from typing import TYPE_CHECKING, Any
from typing_extensions import Self
import numpy as np
import ruamel.yaml as ry
from qcodes.instrument import Instrument
from qcodes.instrument.instrument_meta import InstrumentMeta
from quantify_core.data.handling import get_datadir
if TYPE_CHECKING:
from pathlib import Path
def represent_enum(representer: ry.Representer, data: Enum) -> ry.ScalarNode:
"""Provide a value-based representation for Enum instances in YAML."""
# An enum value can be any scalar type; if there's no matching representer, default to str.
repr_method = representer.yaml_representers.get(type(data.value), str)
return repr_method(representer, data.value)
def represent_ndarray(representer: ry.Representer, data: np.ndarray) -> ry.MappingNode:
"""Represent a NumPy array as a mapping, including its type and shape."""
node = representer.represent_mapping(
"!numpy.ndarray",
{
"dtype": str(data.dtype),
"shape": data.shape,
"data": data.tolist(),
},
)
# Set flow_style for "shape" and "data", so they aren't serialized one element per line.
for key in node.value:
if key[0].value in ("shape", "data"):
key[1].flow_style = True
return node
def construct_ndarray(constructor: ry.Constructor, node: ry.MappingNode) -> np.ndarray:
"""Restore a NumPy array from a mapping with its proper type and shape."""
if isinstance(constructor, ry.RoundTripConstructor):
data = ry.CommentedMap()
constructor.construct_mapping(node, maptyp=data, deep=True)
else:
data = constructor.construct_mapping(node, deep=True)
return np.array(data["data"], dtype=data["dtype"]).reshape(data["shape"])
# The "rt" (round-trip) loader can be advantageous compared to the "safe" loader,
# particularly when working with complex, nested Python classes:
# - Comments: it retains comments present in the YAML file
# - Key Order: it preserves the order of keys in YAML mappings (dictionaries)
# - Formatting: it attempts to maintain the original indentation and whitespace
# - Anchors and Aliases: it preserves YAML anchors and aliases
# - Tags: it handles tags, including custom tags that are used to represent your custom classes
[docs]
yaml = ry.YAML(typ="rt")
# Support Enum and its subclasses
yaml.representer.add_multi_representer(Enum, represent_enum)
# Support NumPy arrays
yaml.representer.add_representer(np.ndarray, represent_ndarray)
yaml.constructor.add_constructor("!numpy.ndarray", construct_ndarray)
class YAMLSerializableMeta(InstrumentMeta):
"""
Metaclass to register mixed in classes with the YAML parser.
Needs to inherit from ``InstrumentMeta`` due to metaclass conflict.
"""
def __init__(cls, name, bases, dct) -> None: # noqa: ANN001, N805
"""
Register the class with the YAML parser so that (de)serialization can happen
automatically.
"""
super().__init__(name, bases, dct)
yaml.register_class(cls)
[docs]
class YAMLSerializable(metaclass=YAMLSerializableMeta):
"""
Mixin to allow (de)serialization of instruments from/to YAML.
NOTE: Only works with ``Instrument`` subclasses, for others use `@yaml.register_class`.
NOTE: `to_yaml` and `from_yaml` methods cannot be created because they would be found
and used by the ``ruamel.yaml`` representers and constructors.
"""
[docs]
def to_yaml_file(
self,
path: str | Path | None = None,
add_timestamp: bool = True,
) -> str:
"""
Convert the object's data structure to a YAML string and store it in a file.
Parameters
----------
path
The path to the directory where the file is created. Default
is `None`, in which case the file will be saved in the directory
determined by :func:`~quantify_core.data.handling.get_datadir()`.
add_timestamp
Specify whether to append timestamp to the filename.
Default is True.
Returns
-------
The name of the file containing the serialized object.
"""
if path is None:
path = get_datadir()
name = getattr(self, "name") # noqa: B009 This is to shut up the linter about self.name
if add_timestamp:
timestamp = datetime.now(timezone.utc).strftime("%Y-%m-%d_%H-%M-%S_%Z")
filename = os.path.join(path, f"{name}_{timestamp}.yaml")
else:
filename = os.path.join(path, f"{name}.yaml")
with open(filename, "w") as file:
yaml.dump(self, file)
return filename
@classmethod
[docs]
def from_yaml_file(cls, filename: str | Path) -> Self:
"""
Read YAML data from a file and convert it to an instance of the attached class.
Parameters
----------
filename
The name of the file containing the serialized object.
Returns
-------
The deserialized object.
"""
with open(filename) as file:
deserialized_obj = yaml.load(file)
return deserialized_obj
def __setstate__(self, state: dict[str, Any]) -> None:
"""
When deserializing an ``Instrument``, add it to the qcodes global registry.
Must be invoked by subclasses if :meth:`__setstate__` is overridden.
"""
if isinstance(self, Instrument):
type(self).record_instance(self)
__all__ = ["YAMLSerializable", "yaml"]
