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

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?

Solution

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.