Ruamel Yaml dump in Python: Need to ensure specific keys are quoted based on level of the data

Viewed 381

Problem Faced: Need a solution that takes into consideration where the data is rather than matching based on keys and quotes or doesn't quote based on that placement when dumping with ruamel.yaml.

Please see the similar question posted before here: JSON to YAML in Python : How to get correct string manipulation? as this is solid context.

Current Data Set:

data2 = {'bar': 'var', 'lar': 'nar'}

data = {'foo': [{
    'zar': 'mar',
    'car': [data2]
}]
}

Current Hacky Solution:

import ruamel.yaml
import yaml
import sys

# Data that is actually ever changing in terms of content, but never structure
data2 = {'bar': 'var', 'lar': 'nar'}

data = {'foo': [{
    'zar': 'mar',
    'car': [data2]
}]
}

# This is the solution portion which doesn't quote if they keys match
#Not a valid solution as the keys are liable to change weekly, but data structure would remain the same

def non_quoted_key(self, data):
    if data == 'foo' or data == 'bar' or data == 'lar':
        data = ruamel.yaml.scalarstring.PlainScalarString(data)
    return self.represent_data(data)

#executing the dump
yaml = ruamel.yaml.YAML()
yaml.default_flow_style = False
yaml.default_style = '"'
yaml.indent = 4
yaml.explicit_start = True
yaml.Representer.represent_key = non_quoted_key
yaml.dump(data, sys.stdout)

Current Output (Which is valid but once again this will fail as soon as the keys change):

---
foo:
-   "zar": "mar"
    "car":
    -   bar: "var"
        lar: "nar"

Issue With Current Solution: The keys are ever changing and new keys are constantly being added, but the data structure will always remain the same. The highest level key of the data structure will always need to be unquoted, keys in the next dictionary set must be quoted, and if those keys contain values that are dictionary sets, those child keys must always be quoted.

Bonus Question: How does one ensure the icons for sequences ([) and mappings ({) are shown in the dump? To give a visual example, using the data above how would you ensure the output always showed sequences and mapping as seen below?

---
foo:
    [
        {
            "car":
                [
                    {
                        bar: "var",
                        lar: "nar"
                    }
                ],
            "zar": "mar"               
        },
    ]
1 Answers

As demonstrated in some other questions, it is often a good idea to start with the output and if ruamel.yaml can reproduce that. Since ruamel.yaml is primarily for preserving comments in block style YAML (and empty lines are continuations of the EOL comments, empty or not, on the previous line), I sprinkeled in a few comments, to see what happens.

import sys
import ruamel.yaml

yaml_str = """---
foo:
    [
        {
            "car":
                [ # c1
                    {
                        bar: "var", # c2

                        lar: "nar", # c3
                    }
                ],
            "zar": "mar"               
        },
    ]
"""

yaml = ruamel.yaml.YAML()
yaml.indent(mapping=4, sequence=4, offset=2)
yaml.preserve_quotes = True
yaml.explicit_start = True
data = yaml.load(yaml_str)
print(data)
yaml.dump(data, sys.stdout)

which gives:

ordereddict([('foo', [ordereddict([('car', [ordereddict([('bar', 'var'), ('lar', 'nar')])]), ('zar', 'mar')])])])
---
foo: [{"car": [{bar: "var", lar: "nar"}], "zar": "mar"}]

From this you may see that getting the quotes should not be a problem, but the flow style with items on seperate lines will be a problem as real comments are ignored in flow style output, and inserting empty comments is not going to help.

When you start with your data and data2, you can recursively walk over the datastructure and convert each key and/or value based on recursion depth. You don't describe explicitly what the rules for non-quoting are, but I assume here the root level keys of the original data structures.

import sys
import ruamel.yaml

DQ = ruamel.yaml.scalarstring.DoubleQuotedScalarString
Pl = ruamel.yaml.scalarstring.PlainScalarString

def set_non_root_strings_double_quoted(d, lvl=0):
    if isinstance(d, dict):
        for k in list(d.keys()):  # need list as d will change
            if not isinstance(k, (DQ, Pl)):
                k = DQ(k) if lvl > 0 else Pl(k)
            d[k] = set_non_root_strings_double_quoted(d.pop(k), lvl=lvl+1)
    elif isinstance(d, list):
        for idx, elem in enumerate(d):
            d[idx] = set_non_root_strings_double_quoted(elem, lvl=lvl+1)
    elif isinstance(d, str):
        d = DQ(d)
    return d

data2 = {'bar': 'var', 'lar': 'nar'}

data = {'foo': [{
    'zar': 'mar',
    'car': [data2]
}]
}

set_non_root_strings_double_quoted(data2)
set_non_root_strings_double_quoted(data)

yaml = ruamel.yaml.YAML()
yaml.explicit_start = True
yaml.indent(mapping=4, sequence=4, offset=2)
yaml.dump(data, sys.stdout)

which gives:

---
foo:
  - "zar": "mar"
    "car":
      - bar: "var"
        lar: "nar"

Since set_non_root_strings_double_quoted works in place you can run it on data2 and data after assigning data2 into data, but the order of doing so important as normal str only get set to DQ or Pl once.

You cannot rely on converting dict/list to CommentedMap/CommentedSeq and assigning comments, but you can use it to convert the output stream to flow style for anything but the root of data (the root of data2 is not set to flow style either, but it is forced when dumping, as you cannot have flow-style withing block-style in YAML).

import sys
import ruamel.yaml

DQ = ruamel.yaml.scalarstring.DoubleQuotedScalarString
Pl = ruamel.yaml.scalarstring.PlainScalarString
Map = ruamel.yaml.CommentedMap
Seq = ruamel.yaml.CommentedSeq

def set_non_root_strings_double_quoted(d, lvl=0):
    if isinstance(d, dict):
        res = Map()
        if lvl == 1: res.fa.set_flow_style()
        for k in d:
            if not isinstance(k, (DQ, Pl)):
                k = DQ(k) if lvl > 0 else Pl(k)
            res[k] = set_non_root_strings_double_quoted(d[k], lvl=lvl+1)
            res.yaml_set_comment_before_after_key(k, after='\n')
    elif isinstance(d, list):
        res = Seq()
        if lvl == 1: res.fa.set_flow_style()
        for elem in d:
            res.append(set_non_root_strings_double_quoted(elem, lvl=lvl+1))
    elif isinstance(d, str):
        res = DQ(d)
    else:
        res = d
    return res

data2 = {'bar': 'var', 'lar': 'nar'}
data2 = set_non_root_strings_double_quoted(data2)

data = {'foo': [{
    'car': [data2],  # <<< in the bonus question car & zar have a different order, so change here
    'zar': 'mar'
}]
}

data = set_non_root_strings_double_quoted(data)

yaml = ruamel.yaml.YAML()
yaml.explicit_start = True
yaml.indent(mapping=4, sequence=4, offset=2)
yaml.dump(data, sys.stdout) # just to show how it looks before postprocessing

which gives:

---
foo: [{"car": [{bar: "var", lar: "nar"}], "zar": "mar"}]

And you can post-process during dumping:

def multiline_flowstyle(s):
    idx = -1

    def peek():
        if idx >= len(s):
            return None
        return s[idx+1]

    res = ""
    lvl = 0
    indent = ' ' * 4
    in_quotes = False  # you could have a bracket or curly brace within a quoted value
    comma_found = False
    while idx < len(s) - 1:
        idx += 1
        ch = s[idx]
        if comma_found:
            comma_found = False
            if ch == ' ':
                res += ',\n' + (lvl * indent) # this gobbels the space
            else: # not sure this ever happens
                res += ',\n' + (lvl * indent) + ch
            continue
        if in_quotes:
            if ch == '"':
                in_quotes = False
            res += ch
            continue
        elif ch == '"':
            in_quotes = True
            res += ch
        elif ch in '{[':
            if peek() == ',':
                idx += 1
                ch += ','
            lvl += 1
            res += '\n' + (lvl * indent) + ch
            if peek() not in '{[':
                lvl += 1
                res += '\n' + (lvl * indent)
        elif ch in '}]':
            if peek() == ',':
                idx += 1
                ch += ','
            lvl -= 1
            res += '\n' + (lvl * indent) + ch
            if peek() not in '}]':
                lvl -= 1
                res += '\n' + (lvl * indent)
        elif ch == ',':
            comma_found = True
        elif ch == '\n':
            res += ch + (lvl * indent)
        else:
            res += ch
    return res

yaml.dump(data, sys.stdout, transform=multiline_flowstyle)

which is close to your requested output:

---
foo: 
    [
        {
            "car": 
                [
                    {
                        bar: "var",
                        lar: "nar"
                    }
                ],
             "zar": "mar"
        }
    ]
Related