[Serializer] #python

Serializer.py

#!/usr/bin/env python3
"""
Library providing convenient classes and methods for writing data to files.
"""
import sys
import json
import pickle

try:
    import yaml
except ImportError:
    yaml = None


class Serializer(object):
    ext = ""
    woptions = ""
    roptions = ""

    @classmethod
    def marshal(cls, input_data):
        raise NotImplementedError()

    @classmethod
    def unmarshal(cls, input_string):
        raise NotImplementedError()


class YAMLSerializer(Serializer):
    ext = "yml"
    woptions = "w"
    roptions = "r"

    @classmethod
    def marshal(cls, input_data):
        return yaml.dump(input_data, default_flow_style=False)

    @classmethod
    def unmarshal(cls, input_string):
        return yaml.load(input_string)


class JSONSerializer(Serializer):
    ext = "json"
    woptions = "w"
    roptions = "r"

    @classmethod
    def marshal(cls, input_data):
        return json.dumps(input_data, indent=2)

    @classmethod
    def unmarshal(cls, input_string):
        return json.loads(input_string)


class PickleSerializer(Serializer):
    ext = "p"
    woptions = "wb"
    roptions = "rb"

    @classmethod
    def marshal(cls, input_data):
        return pickle.dumps(input_data)

    @classmethod
    def unmarshal(cls, input_bytes):
        return pickle.loads(input_bytes)


def get_serializer(serializer):
    if serializer == "json":
        return JSONSerializer
    elif serializer == "pickle":
        return PickleSerializer
    elif serializer == "yaml" and yaml is not None:
        return YAMLSerializer
    elif serializer == "yaml" and yaml is None:
        print("You must have PyYAML installed to use YAML as the serializer.\n"
              "Switching to JSON as the serializer.", file=sys.stderr)
    return JSONSerializer


def get_serializer_from_ext(ext):
    if ext == ".json":
        return JSONSerializer
    elif ext == ".p":
        return PickleSerializer
    elif ext in (".yaml", ".yml") and yaml is not None:
        return YAMLSerializer
    elif ext in (".yaml", ".yml") and yaml is None:
        print("You must have PyYAML installed to use YAML as the serializer.\n"
              "Switching to JSON as the serializer.", file=sys.stderr)
    return JSONSerializer