Source

tutagx / tutagx / meta / to_yaml.py

Full commit
from collections import namedtuple
import yaml

from tutagx.meta import model, encode

# By default, the "scalar style" is all-or-nothing.
# Either *every single scalar* is forced into a particular style (which
# leads to an unholy mess with all the numbers, short strings, ``$id``s, etc)
# or everything's left simple and strings with newlines use ugly quotes.
# We want the "simple" style for most things, but we also want nice >-style
# formatting for longer texts -- hence, this hack.
# Based on http://stackoverflow.com/a/7445560/395760
def _custom_string_representer(dumper, data):
    TAG = 'tag:yaml.org,2002:str'
    if '\n' in data:
        return dumper.represent_scalar(TAG, data, style='>')
    return dumper.represent_scalar(TAG, data, style=None)


yaml.add_representer(str, _custom_string_representer)


WriterContext = namedtuple('WriterContext', 'seen')


class YAMLWriter:
    def __init__(self, model, context=None):
        self._encode = encode.make_encoder(model)
        if context is None:
            context = WriterContext(seen={})
        self.context = context

    def encode(self, obj):
        return self.encode_many([obj])[0]

    def encode_many(self, objs):
        return [self._encode(obj, self.context.seen) for obj in objs]

    def dump(self, obj, f):
        self.dump_many([obj], f)

    def dump_many(self, objs, f):
        vals = self.encode_many(objs)
        yaml.safe_dump_all(
            vals, f,
            default_flow_style=False,
            default_style=None,
            indent=4
        )

    def dumps(self, obj):
        return yaml.safe_dump(self.encode(obj))