[python] PyYAML을 사용하여 YAML 파일의 주석 처리하기

YAML은 사람이 쉽게 읽고 쓸 수 있는 데이터 직렬화 형식입니다. 일반적으로 소스 코드의 설정 파일이나 데이터 저장에 많이 사용됩니다.

PyYAML은 YAML 파일을 파싱하고 생성하는 파이썬 라이브러리입니다. 주석은 YAML 파일에서 추가 문서 설명이나 메모를 작성하는 데 사용됩니다. 하지만 PyYAML은 주석을 파싱하지 않기 때문에 기본적으로 주석이 유지되지 않고 무시됩니다.

그러나 우리는 PyYAML을 약간 수정하여 YAML 파일의 주석을 처리할 수 있습니다. 다음은 PyYAML을 사용하여 YAML 파일의 주석을 처리하는 방법을 보여주는 예제 코드입니다.

import yaml
import re

def add_comments(loader, node):
    if isinstance(node, yaml.SequenceNode):
        # 주석에 해당하는 항목을 찾고
        comments = loader.directives.get_comments(node.start_mark.line)
        if comments:
            # 주석을 확장된 노드에 추가
            extended_node = yaml.SequenceNode(node.tag, node.value, node.start_mark, node.end_mark)
            extended_node.ca.comment = comments
            return extended_node

    elif isinstance(node, yaml.MappingNode):
        # 주석에 해당하는 키 또는 값을 찾고
        for knode, vnode in node.value:
            comments = loader.directives.get_comments(knode.start_mark.line)
            if comments:
                # 주석을 확장된 노드에 추가
                extended_key = yaml.ScalarNode('tag:yaml.org,2002:str', knode.value, knode.start_mark, knode.end_mark)
                extended_key.ca.comment = comments
                node.value.remove((knode, vnode))
                node.value.append((extended_key, vnode))

            if isinstance(vnode, yaml.ScalarNode):
                comments = loader.directives.get_comments(vnode.start_mark.line)
                if comments:
                    # 주석을 확장된 노드에 추가
                    extended_value = yaml.ScalarNode(vnode.tag, vnode.value, vnode.start_mark, vnode.end_mark)
                    extended_value.ca.comment = comments
                    node.value.remove((knode, vnode))
                    node.value.append((knode, extended_value))

    return node

def yaml_load(stream):
    pattern = re.compile(r'^#(\s+.*)$')
    comments = {}
    for index, line in enumerate(stream, start=1):
        match = pattern.match(line)
        if match:
            comments[index] = match.group(1)
    
    loader = yaml.Loader(stream)
    loader.directives.comments = comments
    loader.add_constructor(yaml.resolver.BaseResolver.DEFAULT_SEQUENCE_TAG, add_comments)
    loader.add_constructor(yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG, add_comments)

    try:
        return loader.get_single_data()
    finally:
        loader.dispose()

def yaml_dump(data, stream=None, **kwargs):
    class IndentedEmitter(yaml.emitter.Emitter):
        def expect_block_sequence_item(self):
            self.write_indent()
            line = self.line
            self.write_indicator('-', True, indention=False)
            self.states.append(self.expect_first_block_sequence_item)
            self.expect_node(sequence=True)
            self.write_comment(line)

        def expect_block_mapping_key(self):
            self.write_indent()
            line = self.line
            self.states.append(self.expect_block_mapping_value)
            self.expect_node(mapping=True)
            self.write_comment(line)

        def write_comment(self, line):
            comment = self.comments.get(line)
            if comment:
                self.write(' #{}'.format(comment))
                self.write_indent()
                self.column = 1

    if 'Dumper' not in kwargs:
        kwargs['Dumper'] = yaml.Dumper

    comments = {}
    for node, _ in yaml.emitter.CommentedBaseEmitter(data, **kwargs).generate():
        if isinstance(node, yaml.SequenceNode):
            comments[node.start_mark.line] = node.ca.comment
        elif isinstance(node, yaml.MappingNode):
            for line, value in node.ca.items():
                comments[line] = value.comment

    yaml.emitter.CommentedBaseEmitter.comments = comments
    yaml.emitter.Emitter = IndentedEmitter

    return yaml.dump(data, stream, **kwargs)

위의 코드는 내부적으로 YAML 파일에서 주석을 추출한 다음 주석이 포함된 노드를 생성합니다. 이는 PyYAML의 기능을 확장하여 주석을 유지할 수 있게 합니다.

이제 PyYAML을 사용하여 YAML 파일에서 주석을 읽거나 쓸 수 있습니다. 다음은 예제 코드를 통해 주석을 처리하는 방법을 보여줍니다.

# YAML 파일 읽기
with open('example.yaml', 'r') as file:
    data = yaml_load(file)

# YAML 파일 쓰기
with open('example.yaml', 'w') as file:
    yaml_dump(data, file)

PyYAML을 이용하여 YAML 파일의 주석 처리를 손쉽게 할 수 있습니다. 이를 통해 YAML 파일을 보다 구조화하고 가독성 있는 형태로 유지할 수 있습니다.

참고 자료