在 ruamel.yaml 中保留评论

Question

I'm trying to manipulate some YAML files using ruamel.yaml , notably removing specific keys. This seems to work, but in the process all comment lines and empty lines that follow the key, up to the next key, are also removed. Minimal example:

from ruamel.yaml import YAML
import sys


yaml_str = """\
# Our app configuration

# foo is our main variable
foo: bar

# foz is also important
foz: quz

# And finally we have our optional configs.
# These are not really mandatory, they can be
# considere as "would be nice".
opt1: foqz
"""


yaml = YAML(typ='rt')
data = yaml.load(yaml_str)

data.pop('foz', None)

yaml.dump(data, sys.stdout)

which gives:

# Our app configuration

# foo is our main variable
foo: bar

# foz is also important
opt1: foqz

Is there maybe a way to avoid this and only remove the key itself, and any inline comment?

Answer 1

The ruamel.yaml documentation states

This preservation is normally not broken unless you severely alter the structure of a component (delete a key in a dict, remove list entries).

so this behavior should not have come as a surprise.

What you can do is extend the comment on the key before the one you delete with the comment on the key you are going to delete (you can actually do it afterwards, as the comment is still there, there is just no longer a key while dumping to associate it with).

import sys
import ruamel.yaml

yaml_str = """\
# Our app configuration

# foo is our main variable
foo: bar

# foz is also important
foz: quz

# And finally we have our optional configs.
# These are not really mandatory, they can be
# considere as "would be nice".
opt1: foqz
"""

undefined = object()

def my_pop(self, key, default=undefined):
    if key not in self:
        if default is undefined:
            raise KeyError(key)
        return default
    keys = list(self.keys())
    idx = keys.index(key)
    if key in self.ca.items:
        if idx == 0:
            raise NotImplementedError('cannot handle moving comment when popping the first key', key)
        prev = keys[idx-1]
        # print('prev', prev, self.ca)
        comment = self.ca.items.pop(key)[2]
        if prev in self.ca.items:
            self.ca.items[prev][2].value += comment.value
        else:
            self.ca.items[prev] = self.ca.items.pop(key)
    res = self.__getitem__(key)
    self.__delitem__(key)
    return res

ruamel.yaml.comments.CommentedMap.pop = my_pop
    
yaml = ruamel.yaml.YAML()
data = yaml.load(yaml_str)

data.pop('foz', None)

yaml.dump(data, sys.stdout)

which gives:

# Our app configuration

# foo is our main variable
foo: bar

# foz is also important


# And finally we have our optional configs.
# These are not really mandatory, they can be
# considere as "would be nice".
opt1: foqz

If you need to be able to pop the first key in a commented map, then you need to inspect self.ca and handle its comment attribute, which is somewhat more complicated.

As always when using these kind of internals, you should pin the version of ruamel.yaml that you are using. This internals will change.