meerschaum.utils.yaml

  1#! /usr/bin/env python
  2# -*- coding: utf-8 -*-
  3# vim:fenc=utf-8
  4
  5"""
  6Meerschaum wrapper around YAML libraries.
  7
  8This is so switching between PyYAML and ruamel.yaml is smoother.
  9"""
 10
 11from meerschaum.utils.misc import filter_keywords
 12from meerschaum.utils.packages import attempt_import, all_packages, _import_module
 13from meerschaum.utils.warnings import error
 14from meerschaum.utils.threading import Lock
 15
 16_lib = None
 17### Also supports 'ruamel.yaml'.
 18_import_name = 'yaml'
 19
 20_yaml = None
 21_dumper = None
 22_locks = {
 23    '_lib': Lock(),
 24    '_yaml': Lock(),
 25}
 26
 27__pdoc__ = {
 28    'attempt_import': False,
 29    'error': False,
 30    'all_packages': False,
 31}
 32
 33
 34def _string_presenter(dumper, data: str):
 35    """
 36    Format strings with newlines as blocks.
 37    https://stackoverflow.com/a/33300001/9699829
 38    """
 39    tag_str = 'tag:yaml.org,2002:str'
 40    kw = {}
 41    if len(data.splitlines()) > 1:
 42        kw['style'] = '|'
 43    return dumper.represent_scalar(tag_str, data, **kw)
 44
 45
 46class yaml:
 47    """
 48    Wrapper around `PyYAML` and `ruamel.yaml` so that we may switch between implementations.
 49    """
 50    global _yaml, _lib, _dumper
 51    if _import_name is None:
 52        error("No YAML library declared.")
 53    with _locks['_lib']:
 54        try:
 55            _lib = _import_module(_import_name)
 56        except (ImportError, ModuleNotFoundError):
 57            _lib = attempt_import(_import_name, split=False, lazy=False, install=True)
 58    with _locks['_yaml']:
 59        _yaml = _lib if _import_name != 'ruamel.yaml' else _lib.YAML()
 60        if _import_name != 'ruamel.yaml':
 61            _yaml.add_representer(str, _string_presenter)
 62            _yaml.representer.SafeRepresenter.add_representer(str, _string_presenter)
 63
 64    @staticmethod
 65    def safe_load(*args, **kw):
 66        """
 67        Execute `safe_load` for `PyYAML` and `load` for `ruamel.yaml`.
 68        """
 69        if _import_name == 'ruamel.yaml':
 70            return _yaml.load(*args, **filter_keywords(_yaml.load, **kw))
 71        return _yaml.safe_load(*args, **filter_keywords(_yaml.safe_load, **kw))
 72
 73    @staticmethod
 74    def load(*args, **kw):
 75        """
 76        Execute `yaml.load()`.
 77        Handles the breaking change at `v6.0` of `PyYAML`
 78        (added `yaml.Loader` as a positional argument).
 79        """
 80        packaging_version = attempt_import('packaging.version')
 81        if (
 82            _import_name == 'yaml'
 83            and packaging_version.parse(_yaml.__version__) >= packaging_version.parse('6.0')
 84            and 'Loader' not in kw
 85        ):
 86            kw['Loader'] = _yaml.Loader
 87
 88        return _yaml.load(*args, **filter_keywords(_yaml.load, **kw))
 89
 90    @staticmethod
 91    def dump(data, stream=None, **kw):
 92        """
 93        Dump to a stream. If no stream is provided, return a string instead.
 94        For `ruamel.yaml`, it dumps into a `StringIO` stream and returns `getvalue()`.
 95        """
 96        get_string = False
 97        if stream is None and _import_name == 'ruamel.yaml':
 98            stream = _lib.compat.StringIO()
 99            get_string = True
100
101        if _import_name == 'yaml' and 'Dumper' not in kw:
102            kw['Dumper'] = get_dumper_class()
103
104        result = _yaml.dump(data, stream, **filter_keywords(_yaml.dump, **kw))
105        if get_string:
106            return stream.getvalue()
107        return result
108
109
110def get_dumper_class():
111    """
112    Return the dumper class to use when writing.
113    Only supports `yaml`.
114    """
115    global _dumper
116    if _dumper is not None:
117        return _dumper
118
119    if _import_name != 'yaml':
120        return None
121
122    class CustomDumper(_yaml.Dumper):
123        """
124        Add an extra line break when writing.
125        """
126        def write_line_break(self, data=None):
127            if len(self.indents) == 1:
128                super(CustomDumper, self).write_line_break(data)
129            super(CustomDumper, self).write_line_break(data)
130
131    _dumper = CustomDumper
132    return _dumper