""" pretty serialization
    original from <http://djangosnippets.org/snippets/2397/>
"""
import sys
from io import StringIO
import datetime
import yaml
try:
    from yaml import CSafeLoader as SafeLoader
except ImportError:
    from yaml import SafeLoader
import pytz
from django.core.serializers.base import DeserializationError
from django.utils import six
from django.core.serializers.pyyaml import (
    Serializer as YamlSerializer, DjangoSafeDumper)
from django.core.serializers.python import (
    Deserializer as PythonDeserializer,
)


class Serializer(YamlSerializer):
    """ utf8-friendly dumpdata management command """
    def end_serialization(self):
        yaml.dump(self.objects, self.stream, allow_unicode=True,
                  default_flow_style=False,
                  Dumper=DjangoSafeDumper, **self.options)


def Deserializer(stream_or_string, **options):  # pylint:disable=C0103
    """
    Deserialize a stream or string of YAML data.
    """
    if isinstance(stream_or_string, bytes):
        stream_or_string = stream_or_string.decode('utf-8')
    if isinstance(stream_or_string, six.string_types):
        stream = StringIO(stream_or_string)
    else:
        stream = stream_or_string
    try:  # pylint:disable=R0101
        output = yaml.load(stream, Loader=SafeLoader)
        for a_model in output:
            for key, value in a_model.items():
                if key == 'fields':
                    for vkey, vvalue in value.items():
                        if isinstance(vvalue, datetime.datetime):
                            value[vkey] = vvalue.replace(tzinfo=pytz.utc)
        for obj in PythonDeserializer(output, **options):
            yield obj
    except GeneratorExit:
        raise
    except Exception as except_info:  # pylint:disable=W0703
        # Map to deserializer error
        six.reraise(
            DeserializationError, DeserializationError(
                except_info), sys.exc_info()[2])