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
 65    @staticmethod
 66    def safe_load(*args, **kw):
 67        """
 68        Execute `safe_load` for `PyYAML` and `load` for `ruamel.yaml`.
 69        """
 70        if _import_name == 'ruamel.yaml':
 71            return _yaml.load(*args, **filter_keywords(_yaml.load, **kw))
 72        return _yaml.safe_load(*args, **filter_keywords(_yaml.safe_load, **kw))
 73
 74
 75    @staticmethod
 76    def load(*args, **kw):
 77        """
 78        Execute `yaml.load()`.
 79        Handles the breaking change at `v6.0` of `PyYAML`
 80        (added `yaml.Loader` as a positional argument).
 81        """
 82        packaging_version = attempt_import('packaging.version')
 83        _args = list(args)
 84        if (
 85            _import_name == 'yaml'
 86            and packaging_version.parse(_yaml.__version__) >= packaging_version.parse('6.0')
 87        ):
 88            _args += [_yaml.Loader]
 89        
 90        return _yaml.load(*_args, **filter_keywords(_yaml.load, **kw))
 91
 92
 93    @staticmethod
 94    def dump(data, stream=None, **kw):
 95        """
 96        Dump to a stream. If no stream is provided, return a string instead.
 97        For `ruamel.yaml`, it dumps into a `StringIO` stream and returns `getvalue()`.
 98        """
 99        get_string = False
100        if stream is None and _import_name == 'ruamel.yaml':
101            stream = _lib.compat.StringIO()
102            get_string = True
103
104        if _import_name == 'yaml' and 'Dumper' not in kw:
105            kw['Dumper'] = get_dumper_class()
106
107        result = _yaml.dump(data, stream, **filter_keywords(_yaml.dump, **kw))
108        if get_string:
109            return stream.getvalue()
110        return result
111
112
113def get_dumper_class():
114    """
115    Return the dumper class to use when writing.
116    Only supports `yaml`.
117    """
118    global _dumper
119    if _dumper is not None:
120        return _dumper
121
122    if _import_name != 'yaml':
123        return None
124
125    class CustomDumper(_yaml.Dumper):
126        """
127        Add an extra line break when writing.
128        """
129        def write_line_break(self, data=None):
130            if len(self.indents) == 1:
131                super(CustomDumper, self).write_line_break(data)
132            super(CustomDumper, self).write_line_break(data)
133
134    _dumper = CustomDumper
135    return _dumper