PyYAML: Working with Yaml having unusual constructor tags

Viewed 246

I have a YAML file containing the following data:

%YAML 1.1
%TAG !u! tag:something.com,2019:
--- !u!1337 &1
MainKey:
  B: 0
  Z: 2
  C:
  - a: 0
    e: 0
    c: 0
  - a: 0
    e: 0
    c: 0

To work with this YAML, I created a constructor, and a representer. After loading, then dumping the file, I obtain the following result:

!<tag:something.com,2019:1337>
MainKey:
  B: 0
  C:
  - a: 0
    c: 0
    e: 0
  - a: 1
    c: 1
    e: 1
  Z: 2

I have two questions:

  • First: How can I keep the header, and constructor tag identical to the original version ?
  • Second: How can I make sure the keys are dumped in the same order, instead of being sorted ?

Here is the code I am working for this example, so it can be reproduced:

import copy
import yaml

from io import StringIO
from typing import Dict, Type, TextIO


class Foo:

    tag = "tag:something.com,2019:1337"

    def __init__(self, obj: Dict) -> None:
        self.obj = copy.deepcopy(obj)

    def __eq__(self, other: 'Foo'):
        return self.obj == other.obj

    @staticmethod
    def loader(yaml_loader: yaml.SafeLoader, node: yaml.nodes.MappingNode) -> 'Foo':
        decoded_dictionary = yaml_loader.construct_mapping(node, deep=True)
        return Foo(decoded_dictionary)

    @staticmethod
    def representer(dumper: yaml.SafeDumper, foo: 'Foo') -> yaml.nodes.MappingNode:
        return dumper.represent_mapping(Foo.tag, foo.obj)


YAML_FILE = """%YAML 1.1
%TAG !u! tag:something.com,2019:
--- !u!1337 &1
MainKey:
  B: 0
  Z: 2
  C:
  - a: 0
    e: 0
    c: 0
  - a: 0
    e: 0
    c: 0
"""

EXPECTED_FOO = Foo({'MainKey': {'B': 0, 'Z': 2, 'C': [{'a': 0, 'e': 0, 'c': 0}, {'a': 0, 'e': 0, 'c': 0}]}})


def get_loader() -> Type[yaml.SafeLoader]:
    yaml_loader = yaml.SafeLoader
    yaml_loader.add_constructor(Foo.tag, Foo.loader)

    return yaml_loader


def get_dumper() -> Type[yaml.SafeDumper]:
    dumper = yaml.SafeDumper
    dumper.sort_keys = False
    dumper.add_representer(Foo, Foo.representer)

    return dumper


def load(file_pointer: TextIO) -> Foo:
    return yaml.load(file_pointer, Loader=get_loader())


def dump(foo: Foo, file_pointer: TextIO) -> None:
    yaml.dump(foo, file_pointer, Dumper=get_dumper())


def test_can_load():
    with StringIO(YAML_FILE) as file_pointer:
        result = load(file_pointer)

    assert result == EXPECTED_FOO


def test_can_dump():
    with StringIO() as file_pointer:
        dump(EXPECTED_FOO, file_pointer)
        file_pointer.seek(0)

        assert file_pointer.read() == YAML_FILE


if __name__ == "__main__":
    test_can_load()
    test_can_dump()
2 Answers

You should consider using ruamel.yaml (disclaimer I am the author of that package) for this kind of round-tripping. Apart from preserving comments it also preserves directives if you use the same YAML() instance for loading and dumping.

You can, but don't need to, make constructors and representers (if you do you lose the comments).

import sys
import ruamel.yaml

file_in = Path('input.yaml')


class UnEscape:
    def __init__(self, tags):
        self._revmap = {tag[0] + tag[1:].replace('!', '%21'): tag for tag in tags}
        # print('revmap', self._revmap)

    def __call__(self, s):
        for k, v in self._revmap.items():
            s = s.replace(k, v)
        return s
    
yaml = ruamel.yaml.YAML()
yaml.preserve_quotes = True
data = yaml.load(file_in)
unescape = UnEscape(yaml.tags)
data.anchor.always_dump = True   # only necessary because there is no alias in the example 
data['MainKey']['C'][1]['c'] = 42
yaml.dump(data, sys.stdout, transform=unescape)

which gives:

%YAML 1.1
%TAG !u! tag:something.com,2019:
--- &1 !u!1337
MainKey:
  B: 0
  Z: 2
  C:
  - a: 0
    e: 0
    c: 0
  - a: 0
    e: 0
    c: 42

The transform parameter to dump is necessary to revert the escaping of the ! in the middle of the tag (which normally comes out as !u%211337 which is equivalent, but not nice for a diff).

If the derictive is always the same, you can also transform using the simpler:

def unescape(s):
    return s.replace('!u%21', '!u!')

The tag is preserved (you have to set always dump to dump it if there is no corresponding alias that would force the anchor to be dumped). There is however no control over the order of the anchor and the tag. !u!1337 &1 loads the same as &1 !u!1337 and ruamel.yaml always dumps the latter as that makes it clear you have an anchor on the tagged construct and not a tag on anchored construct.

ruamel.yaml always preserves the order of the keys.

When you are using yaml. You are passing (loading it) as a Foo object.

So pyyaml attempts to sterilise it. Just like the pickle module. Try to pass it as a dict.

SideNote: Be careful when using pyyaml as it is unsafe as pickle. And such can result in unwanted code execution.

So it is best to call yaml.safe_load instead of yaml.load

So. Try this:

import copy
import yaml

from io import StringIO
from typing import Dict, Type, TextIO

YAML_FILE = """
MainKey:
  B: 0
  Z: 2
  C:
  - a: 0
    e: 0
    c: 0
  - a: 0
    e: 0
    c: 0
"""

EXPECTED_FOO = {'MainKey': {'B': 0, 'Z': 2, 'C': [{'a': 0, 'e': 0, 'c': 0}, {'a': 0, 'e': 0, 'c': 0}]}}

def test_can_load():
    with StringIO(YAML_FILE) as file_pointer:
        result = yaml.safe_load(file_pointer)

    assert result == EXPECTED_FOO


def test_can_dump():
    with StringIO() as file_pointer:
        yaml.dump(EXPECTED_FOO, file_pointer)
        file_pointer.seek(0)


if __name__ == "__main__":
    test_can_load()
    test_can_dump()

conclusion: If you want to store data(classed) it is best to use pickle If you want to store json compatable data with human readability then it is best to use yaml

I am sorry If I misunderstood your question.

Related