# -*- coding: utf-8 -*-
"""
This file extends the ruamel.yaml package functionality to load and dump more data types needed by
qudi (mostly numpy array and number types).
Provides easy to use yaml_load and yaml_dump functions to read and write qudi YAML files.
Copyright (c) 2021, the qudi developers. See the AUTHORS.md file at the top-level directory of this
distribution and on <https://github.com/Ulm-IQO/qudi-core/>
This file is part of qudi.
Qudi is free software: you can redistribute it and/or modify it under the terms of
the GNU Lesser General Public License as published by the Free Software Foundation,
either version 3 of the License, or (at your option) any later version.
Qudi is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY;
without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
See the GNU Lesser General Public License for more details.
You should have received a copy of the GNU Lesser General Public License along with qudi.
If not, see <https://www.gnu.org/licenses/>.
"""
__all__ = ['SafeRepresenter', 'SafeConstructor', 'YAML', 'yaml_load', 'yaml_dump', 'ParserError',
'YAMLError', 'MarkedYAMLError', 'YAMLStreamError', 'ScannerError', 'ConstructorError',
'DuplicateKeyError']
import os
import numpy as np
import ruamel.yaml as _yaml
from ruamel.yaml.error import YAMLError, MarkedYAMLError, YAMLStreamError
from ruamel.yaml.parser import ParserError, ScannerError
from ruamel.yaml.constructor import ConstructorError, DuplicateKeyError
from enum import Enum, IntEnum, IntFlag, Flag
from importlib import import_module
from collections import OrderedDict
from io import BytesIO, TextIOWrapper
from typing import Optional, Any, Mapping, Dict, Union
_FilePath = Union[str, bytes, os.PathLike]
[docs]
class SafeRepresenter(_yaml.SafeRepresenter):
"""Custom YAML representer for qudi config files.
"""
ndarray_max_size = 20
[docs]
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self._extndarray_count = 0
def ignore_aliases(self, ignore_data):
"""Ignore aliases and anchors. Overwrites base class implementation.
"""
return True
def represent_numpy_int(self, data):
"""Representer for numpy int scalars.
"""
return self.represent_int(data.item())
def represent_numpy_float(self, data):
"""Representer for numpy float scalars.
"""
return self.represent_float(data.item())
def represent_numpy_complex(self, data):
"""Representer for numpy complex scalars.
"""
return self.represent_complex(data.item())
def represent_dict_no_sort(self, data):
"""Representer for dict and OrderedDict to prevent ruamel.yaml from sorting keys.
"""
return self.represent_dict(data.items())
def represent_complex(self, data):
"""Representer for builtin complex type.
"""
return self.represent_scalar(tag='tag:yaml.org,2002:complex', value=str(data))
def represent_frozenset(self, data):
"""Representer for builtin frozenset type.
"""
node = self.represent_set(data)
node.tag = 'tag:yaml.org,2002:frozenset'
return node
def represent_enum(self, data):
"""Representer for enum types with base class enum.
"""
class_name = data.__class__.__name__
module = data.__class__.__module__
try:
mod = import_module(module)
cls = getattr(mod, class_name)
assert data == cls[data.name]
except (AttributeError, ImportError, AssertionError):
raise TypeError(f'Data can not be represented as enum.Enum.')
return self.represent_scalar(tag='tag:yaml.org,2002:enum',
value=f'{module}.{class_name}[{data.name}]')
def represent_flag(self, data):
"""Representer for enum types with base class enum.
"""
class_name = data.__class__.__name__
module = data.__class__.__module__
try:
mod = import_module(module)
cls = getattr(mod, class_name)
assert data == cls(data.value)
except (AttributeError, ImportError, AssertionError):
raise TypeError(f'Data can not be represented as enum.Flag')
return self.represent_scalar(tag='tag:yaml.org,2002:flag',
value=f'{module}.{class_name}({data.value:d})')
def represent_ndarray(self, data):
"""Representer for numpy.ndarrays.
Will represent the array in binary representation as ASCII-encoded string by default.
If the output stream to dump to is a "regular" open text file handle (io.TextIOWrapper) and
the array size exceeds the specified maximum ndarray size, it is dumped into a separate
binary .npy file and is represented in YAML as file path string.
"""
# Write to separate file if possible and required (array size > self.ndarray_max_size)
# FIXME: Find a better way... this is a mean hack to get the file path to dump,
if isinstance(self.dumper._output, TextIOWrapper) and data.size > self.ndarray_max_size:
try:
out_stream_path = self.dumper._output.name
dir_path = os.path.dirname(out_stream_path)
file_name = os.path.splitext(os.path.basename(out_stream_path))[0]
file_path = f'{os.path.join(dir_path, file_name)}-{self._extndarray_count:06}.npy'
np.save(file_path, data, allow_pickle=False, fix_imports=False)
self._extndarray_count += 1
return self.represent_scalar(tag='tag:yaml.org,2002:extndarray', value=file_path)
except:
pass
# Represent as binary stream (ASCII-encoded) by default
with BytesIO() as f:
np.save(f, data, allow_pickle=False, fix_imports=False)
binary_repr = f.getvalue()
node = self.represent_binary(binary_repr)
node.tag = 'tag:yaml.org,2002:ndarray'
return node
# register custom representers
SafeRepresenter.add_representer(frozenset, SafeRepresenter.represent_frozenset)
SafeRepresenter.add_representer(complex, SafeRepresenter.represent_complex)
SafeRepresenter.add_representer(dict, SafeRepresenter.represent_dict_no_sort)
SafeRepresenter.add_representer(OrderedDict, SafeRepresenter.represent_dict_no_sort)
SafeRepresenter.add_representer(np.ndarray, SafeRepresenter.represent_ndarray)
SafeRepresenter.add_multi_representer(Enum, SafeRepresenter.represent_enum)
SafeRepresenter.add_multi_representer(IntEnum, SafeRepresenter.represent_enum)
SafeRepresenter.add_multi_representer(Flag, SafeRepresenter.represent_flag)
SafeRepresenter.add_multi_representer(IntFlag, SafeRepresenter.represent_flag)
SafeRepresenter.add_multi_representer(np.integer, SafeRepresenter.represent_numpy_int)
SafeRepresenter.add_multi_representer(np.floating, SafeRepresenter.represent_numpy_float)
SafeRepresenter.add_multi_representer(np.complexfloating, SafeRepresenter.represent_numpy_complex)
[docs]
class SafeConstructor(_yaml.SafeConstructor):
"""Custom YAML constructor for qudi config files.
"""
def construct_ndarray(self, node):
"""The constructor for a numpy array that is saved as binary string with ASCII-encoding.
"""
value = self.construct_yaml_binary(node)
with BytesIO(value) as f:
return np.load(f)
def construct_extndarray(self, node):
"""The constructor for a numpy array that is saved in a separate file.
"""
return np.load(self.construct_yaml_str(node), allow_pickle=False, fix_imports=False)
def construct_frozenset(self, node):
"""The frozenset constructor.
"""
try:
# FIXME: The returned generator does not properly work with iteration using next()
return frozenset(tuple(self.construct_yaml_set(node))[0])
except IndexError:
return frozenset()
def construct_complex(self, node):
"""The complex constructor.
"""
return complex(self.construct_yaml_str(node))
def construct_enum(self, node):
"""The Enum constructor.
"""
enum_repr_str = self.construct_yaml_str(node)
enum_mod_cls, enum_name = enum_repr_str.rsplit(']', 1)[0].rsplit('[', 1)
module, cls_name = enum_mod_cls.rsplit('.', 1)
cls = getattr(import_module(module), cls_name)
return cls[enum_name]
def construct_flag(self, node):
"""The Flag constructor.
"""
enum_repr_str = self.construct_yaml_str(node)
enum_mod_cls, enum_value_str = enum_repr_str.rsplit(')', 1)[0].rsplit('(', 1)
module, cls_name = enum_mod_cls.rsplit('.', 1)
cls = getattr(import_module(module), cls_name)
return cls(int(enum_value_str))
# register custom constructors
SafeConstructor.add_constructor('tag:yaml.org,2002:frozenset', SafeConstructor.construct_frozenset)
SafeConstructor.add_constructor('tag:yaml.org,2002:complex', SafeConstructor.construct_complex)
SafeConstructor.add_constructor('tag:yaml.org,2002:ndarray', SafeConstructor.construct_ndarray)
SafeConstructor.add_constructor('tag:yaml.org,2002:extndarray',
SafeConstructor.construct_extndarray)
SafeConstructor.add_constructor('tag:yaml.org,2002:enum', SafeConstructor.construct_enum)
SafeConstructor.add_constructor('tag:yaml.org,2002:flag', SafeConstructor.construct_flag)
[docs]
class YAML(_yaml.YAML):
"""ruamel.yaml.YAML subclass to be used by qudi for all loading/dumping purposes.
Will always use the 'safe' option without round-trip functionality.
"""
[docs]
def __init__(self, **kwargs):
"""
Parameters
----------
kwargs
Keyword arguments accepted by ruamel.yaml.YAML(), excluding "typ".
"""
kwargs['typ'] = 'safe'
super().__init__(**kwargs)
self.default_flow_style = False
self.Representer = SafeRepresenter
self.Constructor = SafeConstructor
[docs]
def yaml_load(file_path: _FilePath, ignore_missing: Optional[bool] = False) -> Dict[str, Any]:
"""Loads a qudi style YAML file.
Raises OSError if the file does not exist or can not be accessed.
Parameters
----------
file_path : str
Path to config file.
ignore_missing : bool, optional
Flag to suppress FileNotFoundError.
Returns
-------
dict
The data as python/numpy objects in a dict.
"""
try:
with open(file_path, 'r') as f:
data = YAML().load(f)
# yaml returns None if the stream was empty
return dict() if data is None else data
except OSError:
if ignore_missing:
return dict()
else:
raise
[docs]
def yaml_dump(file_path: _FilePath, data: Mapping[str, Any]) -> None:
"""Saves data to file_path in qudi style YAML format. Creates subdirectories if needed.
file_path : str
Path to YAML file to save data into.
data : dict
Dict containing the data to save to file.
"""
file_dir = os.path.dirname(file_path)
if file_dir:
os.makedirs(file_dir, exist_ok=True)
with open(file_path, 'w') as f:
YAML().dump(data, f)