How do I align the eol comments in ruamel.yaml so that they are all in the same column

Viewed 146

I am currently generating YAML files automatically with some configuration values. Sometimes these values have a comment attached and the comment is added to the YAML file.

Simple example

import sys
from ruamel.yaml import CommentedMap, YAML

top = CommentedMap()
top['sub_key1'] = data = CommentedMap()

data['a'] = 1
data['b'] = 'asdf'
data['c'] = 3.3333
data.yaml_add_eol_comment('comment 1', 'a')
data.yaml_add_eol_comment('comment 2', 'b')
data.yaml_add_eol_comment('comment 3', 'c')

top['sub_key2'] = data = CommentedMap()

data['a'] = 'long text'
data['b'] = 'an even longer text'
data.yaml_add_eol_comment('comment 4', 'a')
data.yaml_add_eol_comment('comment 5', 'b')

YAML().dump(top, sys.stdout)

This works and outputs this as expected

sub_key1:
  a: 1  # comment 1
  b: asdf # comment 2
  c: 3.3333 # comment 3
sub_key2:
  a: long text  # comment 4
  b: an even longer text # comment 5

However I'd really like to have the comments aligned like this (or even better with two spaces after the value).

sub_key1:
  a: 1      # comment 1
  b: asdf   # comment 2
  c: 3.3333 # comment 3
sub_key2:
  a: long text           # comment 4
  b: an even longer text # comment 5

I can't use the column parameter when adding the eol comment because the column is absolute from the start and I don't know

  • on which level I am
  • how long the generated key/value pair is
  • what the current indentation i

Is there any way to align the comments after creation?

2 Answers

You could dump without setting the column, read back the resulting YAML and analyse the comment structure ( data['sub_key1']ca.items ) to get what is the largest column value and then dump one more time using that value.

Or while you have the data loaded with comments, just set all column values on the comment to the maximum or the maximum plus one to get two spaces after the longest value:

import sys
import io
from ruamel.yaml import CommentedMap, YAML

yaml = YAML()
top = CommentedMap()
top['sub_key1'] = data = CommentedMap()

data['a'] = 1
data['b'] = 'asdf'
data['c'] = 3.3333
data.yaml_add_eol_comment('comment 1', 'a')
data.yaml_add_eol_comment('comment 2', 'b')
data.yaml_add_eol_comment('comment 3', 'c')

top['sub_key2'] = data = CommentedMap()

data['a'] = 'long text'
data['b'] = 'an even longer text'
data.yaml_add_eol_comment('comment 4', 'a')
data.yaml_add_eol_comment('comment 5', 'b')

buf = io.BytesIO()
yaml.dump(top, buf)

top = yaml.load(buf.getvalue())

def align_comments(d, extra=0):
    def align_one(d, extra=0):
        comments = d.ca.items.values()
        if not comments:
            return
        max = -1
        for comment in comments:
            if comment[2].column > max:
                max = comment[2].column
        for comment in comments:
            comment[2].column = max + extra

    if isinstance(d, dict):
        align_one(d, extra=extra)
        for val in d.values():
            align_comments(val, extra=extra)
    elif isinstance(d, list):
        align_one(d, extra=extra)
        for elem in d:
            align_comments(elem, extra=extra)


align_comments(top, extra=1)

yaml.dump(top, sys.stdout)

which gives:

sub_key1:
  a: 1       # comment 1
  b: asdf    # comment 2
  c: 3.3333  # comment 3
sub_key2:
  a: long text            # comment 4
  b: an even longer text  # comment 5

Make sure you pin the version of ruamel.yaml that you are using. As these kind of internals will change.

@Anthon: I modified your example a little bit for better readability. Thank you for your help

def align_comments(d, extra_indent=0):

    is_dict = isinstance(d, dict)
    if not is_dict and not isinstance(d, list):
        return None

    comments = d.ca.items.values()
    if comments:
        max_col = max(map(lambda x: x[2].column, comments), default=0)
        for comment in comments:
            comment[2].column = max_col + extra_indent

    for element in (d.values() if is_dict else d):
        align_comments(element, extra_indent=extra_indent)
    return None

Usage:

buf = io.BytesIO()
yaml.dump(top, buf)

top = yaml.load(buf.getvalue())
align_comments(top, extra=1)

yaml.dump(top, sys.stdout)
Related