Chainerable JSONEncoder in Python

JSONEncoderChain.py

from __future__ import absolute_import

import json
from abc import ABCMeta, abstractmethod
from datetime import datetime
from json import JSONEncoder as JSONEncoder

from six import add_metaclass


class JSONEncoderChain(JSONEncoder):
    def __init__(self, skipkeys=False, ensure_ascii=True,
                 check_circular=True, allow_nan=True, sort_keys=False,
                 indent=None, separators=None, default=None):
        super().__init__(skipkeys, ensure_ascii, check_circular, allow_nan, sort_keys, indent, separators, default)
        self._encoders = []

    def add_encoder(self, encoder):
        if not isinstance(encoder, (tuple, AbstractJSONEncoderWithChecker)):
            raise TypeError("encoder must be either (chk_f, enc_f) or AbstractEncoderWithChecker")
        self._encoders.append(encoder)
        return self

    def default(self, o):
        for encoder in self._encoders:
            if isinstance(encoder, tuple):
                chk_f, enc_f = encoder
            elif isinstance(encoder, AbstractJSONEncoderWithChecker):
                chk_f, enc_f = encoder.check, encoder.encode
            else:
                continue
            if chk_f(o):
                return enc_f(o)
        else:
            return JSONEncoder.default(self, o)

    def dumps(self, o, **kwargs):
        return json.dumps(o, default=self.default, **kwargs)

    def dump(self, o, fp, **kwargs):
        return json.dump(o, fp, default=self.default, **kwargs)


@add_metaclass(ABCMeta)
class AbstractJSONEncoderWithChecker(JSONEncoder):
    @abstractmethod
    def check(self, o):
        pass

    @abstractmethod
    def encode(self, o):
        pass

    def default(self, o):
        return self.encode(o) if self.check(o) else JSONEncoder.default(self, o)


class DateTimeJSONEncoder(AbstractJSONEncoderWithChecker):
    def check(self, o):
        return isinstance(o, datetime)

    def encode(self, o):
        return o.isoformat()