Generating Python type annotations from JSON

json2types.py

import typing


def convert(it, context: str):
    if isinstance(it, dict):
        return typing.TypedDict(
            context,
            {
                k: convert(v, f"{context.capitalize()}{k.capitalize()}")
                for (k, v) in it.items()
            },
        )
    elif isinstance(it, str):
        return str
    elif isinstance(it, int):
        return int
    elif isinstance(it, float):
        return float
    elif it is None:
        return typing.Any
    else:
        raise NotImplementedError(type(it))


def nested(typedef):
    try:
        annotations = typedef.__annotations__
    except AttributeError:
        pass
    else:
        for k, v in annotations.items():
            # noinspection PyUnresolvedReferences,PyProtectedMember
            if isinstance(v, typing._TypedDictMeta):
                yield from nested(v)
                yield v


def codegen_typeddict(
    it,
    default_name: str,
    pycharm: bool,
    include_undefined: bool,
):
    s = []
    if pycharm:
        s.append("# noinspection DuplicatedCode")

    classname = it.__name__ or default_name
    s.append(f"class {classname}(TypedDict):")

    fields = []
    for k, v in it.__annotations__.items():
        try:
            v = v.__name__
        except AttributeError:
            v = str(v)
        if v == "typing.Any" and not include_undefined:
            continue
        v = v.replace("typing.", "")
        fields.append(f"    {k}: {v}")

    if not fields:
        fields = ["    pass"]

    s.extend(fields)

    return "\n".join(s)


def codegen(it, name, pycharm, include_undefined):
    typedef = convert(it, context="")
    s = []

    imports = ["TypedDict"]
    if include_undefined:
        imports.append("Any")

    s.append(f"from typing import {', '.join(sorted(imports))}")
    s.append("")
    s.append("")

    for d in nested(typedef):
        s.append(codegen_typeddict(d, name, pycharm, include_undefined))
        s.append("")
        s.append("")
    s.append(codegen_typeddict(typedef, name, pycharm, include_undefined))
    return "\n".join(s)


def cli():
    import argparse
    import json
    import sys

    parser = argparse.ArgumentParser()
    parser.add_argument(
        "input",
        type=argparse.FileType("r", encoding="utf-8"),
    )
    parser.add_argument(
        "-o",
        "--output",
        type=argparse.FileType("w", encoding="utf-8"),
        default=sys.stdout,
    )
    parser.add_argument(
        "-n",
        "--name",
        default="Root",
        help="Name for the top-level type",
    )
    parser.add_argument(
        "-p",
        "--pycharm",
        action="store_true",
        help="Include PyCharm metadata",
    )
    parser.add_argument(
        "-u",
        "--undefined",
        action="store_true",
        help="Include type annotations for undefined values",
    )
    args = parser.parse_args()

    args.output.write(f'"""\nThis file was generated using {__file__}\n"""\n\n')
    args.output.write(
        codegen(
            json.load(args.input),
            args.name,
            args.pycharm,
            args.undefined,
        )
    )
    args.output.write("\n")


if __name__ == "__main__":
    cli()