1"""
2    salt.utils.yamldumper
3    ~~~~~~~~~~~~~~~~~~~~~
4
5"""
6# pylint: disable=W0232
7#         class has no __init__ method
8
9
10import collections
11
12import salt.utils.context
13import yaml  # pylint: disable=blacklisted-import
14from salt.utils.odict import OrderedDict
15
16try:
17    from yaml import CDumper as Dumper
18    from yaml import CSafeDumper as SafeDumper
19except ImportError:
20    from yaml import Dumper
21    from yaml import SafeDumper
22
23
24__all__ = [
25    "OrderedDumper",
26    "SafeOrderedDumper",
27    "IndentedSafeOrderedDumper",
28    "get_dumper",
29    "dump",
30    "safe_dump",
31]
32
33
34class IndentMixin(Dumper):
35    """
36    Mixin that improves YAML dumped list readability
37    by indenting them by two spaces,
38    instead of being flush with the key they are under.
39    """
40
41    def increase_indent(self, flow=False, indentless=False):
42        return super().increase_indent(flow, False)
43
44
45class OrderedDumper(Dumper):
46    """
47    A YAML dumper that represents python OrderedDict as simple YAML map.
48    """
49
50
51class SafeOrderedDumper(SafeDumper):
52    """
53    A YAML safe dumper that represents python OrderedDict as simple YAML map.
54    """
55
56
57class IndentedSafeOrderedDumper(IndentMixin, SafeOrderedDumper):
58    """
59    A YAML safe dumper that represents python OrderedDict as simple YAML map,
60    and also indents lists by two spaces.
61    """
62
63
64def represent_ordereddict(dumper, data):
65    return dumper.represent_dict(list(data.items()))
66
67
68def represent_undefined(dumper, data):
69    return dumper.represent_scalar("tag:yaml.org,2002:null", "NULL")
70
71
72OrderedDumper.add_representer(OrderedDict, represent_ordereddict)
73SafeOrderedDumper.add_representer(OrderedDict, represent_ordereddict)
74SafeOrderedDumper.add_representer(None, represent_undefined)
75
76OrderedDumper.add_representer(
77    collections.defaultdict, yaml.representer.SafeRepresenter.represent_dict
78)
79SafeOrderedDumper.add_representer(
80    collections.defaultdict, yaml.representer.SafeRepresenter.represent_dict
81)
82OrderedDumper.add_representer(
83    salt.utils.context.NamespacedDictWrapper,
84    yaml.representer.SafeRepresenter.represent_dict,
85)
86SafeOrderedDumper.add_representer(
87    salt.utils.context.NamespacedDictWrapper,
88    yaml.representer.SafeRepresenter.represent_dict,
89)
90
91OrderedDumper.add_representer(
92    "tag:yaml.org,2002:timestamp", OrderedDumper.represent_scalar
93)
94SafeOrderedDumper.add_representer(
95    "tag:yaml.org,2002:timestamp", SafeOrderedDumper.represent_scalar
96)
97
98
99def get_dumper(dumper_name):
100    return {
101        "OrderedDumper": OrderedDumper,
102        "SafeOrderedDumper": SafeOrderedDumper,
103        "IndentedSafeOrderedDumper": IndentedSafeOrderedDumper,
104    }.get(dumper_name)
105
106
107def dump(data, stream=None, **kwargs):
108    """
109    .. versionadded:: 2018.3.0
110
111    Helper that wraps yaml.dump and ensures that we encode unicode strings
112    unless explicitly told not to.
113    """
114    if "allow_unicode" not in kwargs:
115        kwargs["allow_unicode"] = True
116    kwargs.setdefault("default_flow_style", None)
117    return yaml.dump(data, stream, **kwargs)
118
119
120def safe_dump(data, stream=None, **kwargs):
121    """
122    Use a custom dumper to ensure that defaultdict and OrderedDict are
123    represented properly. Ensure that unicode strings are encoded unless
124    explicitly told not to.
125    """
126    if "allow_unicode" not in kwargs:
127        kwargs["allow_unicode"] = True
128    kwargs.setdefault("default_flow_style", None)
129    return yaml.dump(data, stream, Dumper=SafeOrderedDumper, **kwargs)
130