Как сохранять комментарии в yaml после его модификации кодом

Question

Мне нужно имея один yaml файл получить из него аналогичный, заменяя некоторые значения в нем (заданы мною в другой файле. Имеется такой скрипт:

import yaml

def load_yaml(file_path):
    with open(file_path, 'r') as file:
        return yaml.safe_load(file)

def modify_keys(data, modifications):
    if isinstance(data, dict):
        for key, value in data.items():
            if key in modifications:
                data[key] = modifications[key]
            else:
                modify_keys(value, modifications)
    elif isinstance(data, list):
        for idx, item in enumerate(data):
            modify_keys(item, modifications)

def save_yaml_with_comments(file_path, data, modifications):
    with open(file_path, 'w') as file:
        yaml_string = yaml.dump(data)
        for line in yaml_string.splitlines():
            if line.strip().startswith("#"):
                file.write(line + "\n")
            else:
                key = line.split(':')[0].strip()
                if key in modifications:
                    new_line = f"{key}: {modifications[key]}"
                    file.write(new_line + "\n")
                else:
                    file.write(line + "\n")

def main(input_file, output_file, modifications_file):
    data = load_yaml(input_file)
    modifications = load_yaml(modifications_file)
    modify_keys(data, modifications)
    save_yaml_with_comments(output_file, data, modifications)

if __name__ == "__main__":
    input_file = 'file1.yaml'         # Исходный YAML файл
    output_file = 'file2.yaml'       # Файл для сохранения изменений
    modifications_file = 'replacements.yaml'  # Файл с ключами и их новыми значениями

    main(input_file, output_file, modifications_file)

file1.yaml имеет состав

name:
#comment
  user1: Igor
  user2: Nikita

replacements.yaml имеет состав

  user1: Vasya

Итог всего будет такой:

name:
  user1: Vasya
  user2: Nikita

Он затирает комментарий, как мне сделать так, чтобы комментарий также сохранялся?

Answer 1

Суть в том, что уже в функции def load_yaml(file_path) у Вас затираются комментарии при использовании yaml.safe_load(file).

Как вариант, можно использовать ruamel.yaml - при использовании обычного (небезопасного) load комментарии сохраняются:

import ruamel.yaml

def load_yaml(yaml_cls, file_path):
    with open(file_path, 'r') as file:
        return yaml_cls.load(file)

def modify_keys(data, modifications):
    if isinstance(data, dict):
        for key, value in data.items():
            if key in modifications:
                data[key] = modifications[key]
            else:
                modify_keys(value, modifications)
    elif isinstance(data, list):
        for idx, item in enumerate(data):
            modify_keys(item, modifications)

def save_yaml(yaml_cls, data, file_path):
    with open(file_path, 'w') as file:
        yaml_cls.dump(data, file)

def main(input_file, output_file, modifications_file):
    yaml = ruamel.yaml.YAML()
    data = load_yaml(yaml, input_file)
    modifications = load_yaml(yaml, modifications_file)
    modify_keys(data, modifications)
    save_yaml(yaml, data, output_file)

if __name__ == "__main__":
    input_file = 'file1.yaml'                 # Исходный YAML файл
    output_file = 'file2.yaml'                # Файл для сохранения изменений
    modifications_file = 'replacements.yaml'  # Файл с ключами и их новыми значениями
    main(input_file, output_file, modifications_file)

P.S.: Код работает, если записи в replacements.yaml идут в "плоском" формате, без - перед значениями:

user1: Vasya
user2: Masha

БЛОГ НА HUSL

Как сохранять комментарии в yaml после его модификации кодом

Ответы (1 шт):