Typed JSON serializability in Python

serializable.py

from __future__ import annotations

import json
from typing import Any, Dict, List, Tuple, Union, cast

from typing_extensions import Protocol, runtime

# Type has to be ignored until mypy support recursive types (should be
# soon-ish, see mypy#731). This means that type errors one step down in a dict
# hierarchy will not be detected, since the "recursed" type will be interpreted
# as Any by mypy, e.g. {True: False} will yield an error but {'foo': {True:
# False}} will not, since the inferred type here will be approximately
# Dict[str, Any].
JSONSerializable = Union[  # type: ignore
    List['JSONSerializable'],
    Tuple['JSONSerializable'],
    Dict[str, 'JSONSerializable'],
    str, int, float, bool, None,
    'CustomJSONSerializable',
]


@runtime
class CustomJSONSerializable(Protocol):
    def json_serializable(self) -> JSONSerializable: ...


class CustomEncoder(json.JSONEncoder):
    def default(self, o: Any) -> JSONSerializable:
        if isinstance(o, CustomJSONSerializable):
            return o.json_serializable()
        return cast(
            JSONSerializable,
            json.JSONEncoder.default(self, o))

usage.py

from .serializable import JSONSerializable, CustomEncoder


class Impl:
    # since Dict[str, str] is a subtype of JSONSerializable, this method
    # implements the CustomJSONSerializable Protocol
    def json_serializable(self) -> Dict[str, str]:
        return {'foo': 'bar'}


class Expl:
    def json_serializable(self) -> JSONSerializable:
        return {'boo': 'far', 'lar': Impl(), }


complex_data: JSONSerializable = {
    'stuff': 123,
    'deeper': {
        'i': Impl(),
        'e': Expl(),
    },
    'da': 'bass',
}

print(json.dumps(complex_data, cls=CustomEncoder, indent=4))

# this raises an error because boolean is not allowed as key for serializable
# dicts
error: JSONSerializable = {True: False}