yaml_schematics.py

from schematics.models import Model
from schematics.types import BaseType, StringType, BooleanType, IntType
from schematics.types.compound import ListType, ModelType as _ModelType, DictType
from schematics.transforms import blacklist, export_loop

import ruamel.yaml as yaml
from ruamel.yaml.comments import CommentedMap
from ruamel.yaml.scalarstring import PreservedScalarString
from ruamel.yaml.compat import string_types, ordereddict


def com_insert(self, pos, key, value, comment=None):
    od = ordereddict()
    od.update(self)
    for k in od:
        del self[k]
    for index, old_key in enumerate(od):
        if pos == index:
            self[key] = value
        self[old_key] = od[old_key]
    if comment is not None:
        self.yaml_add_eol_comment(comment, key=key)

CommentedMap.insert = com_insert


class YamlModel(Model):
    """
    A Schematics Model that can load, edit and dump yaml, keeping formatting
    and comments.
    """

    def __init__(self, raw_data=None, **kwargs):
        self._raw_data = raw_data

        if not raw_data:
            raw_data = {}

        super().__init__(raw_data, **kwargs)

        self._bind_fields()

    def import_data(self, raw_data, **kwargs):
        new = False
        if not self._raw_data:
            self._raw_data = raw_data
            new = True

        data = super().import_data(raw_data, **kwargs)

        if new:
            self._bind_fields()

        return data

    def _bind_fields(self):
        if isinstance(self._raw_data, CommentedMap):
            for k, field in self._fields.items():
                try:
                    field.orig_value = self._raw_data[k]
                except KeyError:
                    field.orig_value = None

    def to_primitive(self, role=None, context=None):
        field_converter = lambda field, value: field.to_primitive(value,
                                                              context=context)
        data = export_loop(self.__class__, self, field_converter, role=role,
                           raise_error_on_role=True)

        return self._process_data(data)

    def to_yaml(self):
        return yaml.round_trip_dump(self.serialize(), width=1000)

    def export_data(self):
        return self.to_primitive(context={'export': True})

    def _process_data(self, data):
        model_keys = [k for k in self._fields.keys() if data.get(k, None) != None]

        def format_value(value):
            if isinstance(value, string_types):
                if '\n' in value:
                    # Example:
                    # {"knight": "multiple\nlines"} should get dumped as:
                    # knight: |-
                    #   multiple
                    #   lines
                    string = value.replace('\r\n', '\n').replace('\r', '\n')
                    return PreservedScalarString(string.strip())
                return value.strip()

            return value

        if isinstance(self._raw_data, CommentedMap):
            cmap = self._raw_data
            for k in list(cmap.keys()):
                if k in model_keys:
                    if isinstance(data[k], bool):
                        # Make sure bool won't get dumped as string (with quotes)
                        cmap[k] = data[k]
                    else:
                        # Type could be anything from int or str to PreservedScalarString
                        cmap[k] = type(cmap[k])(data[k])
                else:
                    del cmap[k]
            for index, key in enumerate(model_keys):
                if cmap.get(key, None) == None:
                    cmap.insert(index, key, format_value(data[key]))
        else:
            cmap = CommentedMap([(k, format_value(data[k])) for k in model_keys])

        return cmap

    class Options:
        serialize_when_none = False


class ModelType(_ModelType):
    """
    Extend the default ModelType with our custom `_process_data` function.
    """

    def export_loop(self, model_instance, field_converter, role=None,
                    print_none=False):
        data = super().export_loop(model_instance, field_converter, role,
                                   print_none)

        return model_instance._process_data(data)


class BoolType(BooleanType):
    """
    String values such as Yes/No and On/Off have been depricated in YAML 1.2
    The purpose of this class is to provide support for the old standard,
    while also converting to true/false upon saving.
    """

    TRUE_VALUES  = ('y', 'Y', 'yes', 'Yes', 'YES',
                    'true', 'True', 'TRUE', True,
                    'on', 'On', 'ON')
    FALSE_VALUES = ('n', 'N', 'no', 'No', 'NO',
                    'false', 'False', 'FALSE', False,
                    'off', 'Off', 'OFF')

If you want to be able to load 1.1 style booleans, it might be more easy to either add version=(1,1) to round_trip_load (or specify %YAML 1.1 at the top of the file), although that also reverts the loader to old-style octals.