Easy structured output types

output.py

import enum
import functools
import json
import uuid
from datetime import datetime, timedelta

import pyrfc3339
import ruamel.yaml


class JSONEcoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, uuid.UUID):
            return str(obj)
        elif isinstance(obj, datetime):
            return pyrfc3339.generate(obj)
        elif isinstance(obj, timedelta):
            return obj.total_seconds()
        return super().default(obj)


class OutputWriter:
    def __init__(self, output_function):
        self._output_function = output_function

    def dump(self, obj, fp):
        self._output_function(obj, fp)


def _repr_enum(representer, data: enum.Enum):
    return representer.represent_data(data.value)


def _repr_uuid(representer, data: uuid.UUID):
    return representer.represent_data(str(data))


def _repr_datetime(representer, data: datetime):
    return representer.represent_data(pyrfc3339.generate(data))


def _repr_timedelta(representer, data: timedelta):
    return representer.represent_data(data.total_seconds())


def _repr_tuple(representer, data):
    return representer.represent_data(tuple(data))


def yaml_writer():
    yaml = ruamel.yaml.YAML()
    yaml.default_flow_style = False

    yaml.representer.add_multi_representer(enum.Enum, _repr_enum)
    yaml.representer.add_multi_representer(uuid.UUID, _repr_uuid)
    yaml.representer.add_multi_representer(datetime, _repr_datetime)
    yaml.representer.add_multi_representer(timedelta, _repr_timedelta)
    yaml.representer.add_multi_representer(tuple, _repr_tuple)

    return OutputWriter(yaml.dump)


def json_writer(cls=None):
    return OutputWriter(functools.partial(json.dump, cls=cls, indent=True))


def json_compact_writer(cls=None):
    return OutputWriter(functools.partial(json.dump, cls=cls))