rhodecode/lib/ext_json.py

import datetime
import decimal
import functools
# we keep simplejson for having dump functionality still
# orjson doesn't support it
import simplejson as sjson

import orjson
import orjson as json


from rhodecode.lib.datelib import is_aware
from rhodecode.lib.str_utils import safe_str

try:
    import rhodecode.translation
except ImportError:
    rhodecode = None

__all__ = ['json']


def _obj_dump(obj):
    """
    Custom function for dumping objects to JSON, if obj has __json__ attribute
    or method defined it will be used for serialization

    :param obj:
    """
    # See "Date Time String Format" in the ECMA-262 specification.
    # some code borrowed from django 1.4
    if isinstance(obj, set):
        return list(obj)
    elif isinstance(obj, datetime.datetime):
        r = obj.isoformat()
        if isinstance(obj.microsecond, int):
            r = r[:23] + r[26:]
        if r.endswith('+00:00'):
            r = r[:-6] + 'Z'
        return r
    elif isinstance(obj, datetime.date):
        return obj.isoformat()
    elif isinstance(obj, datetime.time):
        if is_aware(obj):
            raise TypeError("Time-zone aware times are not JSON serializable")
        r = obj.isoformat()
        if isinstance(obj.microsecond, int):
            r = r[:12]
        return r
    elif hasattr(obj, '__json__'):
        if callable(obj.__json__):
            return obj.__json__()
        else:
            return obj.__json__
    elif isinstance(obj, decimal.Decimal):
        return str(obj)
    elif isinstance(obj, complex):
        return [obj.real, obj.imag]
    elif rhodecode and isinstance(obj, rhodecode.translation._LazyString):
        return obj.eval()
    else:
        raise TypeError(repr(obj) + " is not JSON serializable")


sjson.dumps = functools.partial(sjson.dumps, default=_obj_dump)
sjson.dump = functools.partial(sjson.dump, default=_obj_dump)

json.dumps = functools.partial(json.dumps, default=_obj_dump, option=orjson.OPT_NON_STR_KEYS)
json.dump = functools.partial(sjson.dump, default=_obj_dump)


def formatted_json(*args, **kwargs):
    # alias for formatted json
    opts = orjson.OPT_NON_STR_KEYS | orjson.OPT_INDENT_2 | orjson.OPT_SORT_KEYS
    return functools.partial(json.dumps, option=opts)(*args, **kwargs)


def formatted_str_json(*args, **kwargs):
    opts = orjson.OPT_NON_STR_KEYS | orjson.OPT_INDENT_2 | orjson.OPT_SORT_KEYS
    closure = functools.partial(json.dumps, option=opts)
    return safe_str(closure(*args, **kwargs))


def str_json(*args, **kwargs):
    closure = functools.partial(json.dumps)
    return safe_str(closure(*args, **kwargs))