json_pretty_print.py

import sys
import json
import argparse
from colored import fg, bg, attr

ARGS = None  # initialized in main.
MAX_WIDTH = 100  # overriden in main.
LENGTH_CACHE = {}
LENGTH_CACHE_HITS = 0
LENGTH_CACHE_MISSES = 0


class Color:
    @staticmethod
    def key(inp: str) -> str:
        if ARGS.nocolor:
            return inp
        return f'{fg("light_blue")}{attr("bold")}{inp}{attr("reset")}'

    @staticmethod
    def strval(inp) -> str:
        if ARGS.nocolor:
            return inp
        return f'{fg("green")}{inp}{attr("reset")}'

    @staticmethod
    def null() -> str:
        if ARGS.nocolor:
            return "null"
        return f'{fg("dark_gray")}{attr("bold")}null{attr("reset")}'


# Returns the length of the given object when represented as a one line string.
# It caches the length of the given object and hence does not work if the object is mutated.
def lenDct(inp):
    global LENGTH_CACHE_HITS
    global LENGTH_CACHE_MISSES

    if ARGS.nocache:
        LENGTH_CACHE_MISSES += 1
        return lenDctInternal(inp)

    ret = LENGTH_CACHE.get(id(inp))
    if ret is None:
        ret = lenDctInternal(inp)
        LENGTH_CACHE[id(inp)] = ret
        LENGTH_CACHE_MISSES += 1
    else:
        LENGTH_CACHE_HITS += 1
    return ret


def lenDctInternal(inp):
    L_COLON = 1
    L_SPACE = 1
    L_COMMA = 1
    L_BRACE = 1

    if type(inp) == dict:
        # {<k1v1>, <k2v2>}
        ret = L_BRACE + (len(inp) - 1) * (L_COMMA + L_SPACE) + L_BRACE
        for k, v in inp.items():
            # '<key>: <value>'
            ret += (lenStr(k) + L_COLON + L_SPACE + lenDct(v))
        return ret
    elif type(inp) == list:
        # [i1, i2]
        ret = L_BRACE + (len(inp) - 1) * (L_COMMA + L_SPACE) + L_BRACE
        for item in inp:
            ret += lenDct(item)
        return ret
    elif type(inp) == str:
        return lenStr(inp)
    elif inp == True:
        return 4
    elif inp == False:
        return 5
    elif type(inp) in {int, float}:
        return len(str(inp))
    elif inp is None:
        return 3  # null
    else:
        print(f'lenDctInternal: Not sure how to handle {inp}')
        sys.exit(1)


def lenStr(inp):
    # '<inp>'
    return 1 + len(inp) + 1


def isDctOrLst(inp):
    return type(inp) in (list, dict)


def wrapQuote(inp: str) -> str:
    return f'"{inp}"'


# Helper to printOneLine. This returns a list of strings when stitched together will be the string
# representation of the given object.
# Example:
# input = {"k1": "v2", "k2": "v2"}
# output can be ['{', '"k1": ', '"v2"', ', ', '"k2": ', '"v2"', ', ', '}']
def oneLineStringParts(inp):
    if type(inp) == dict:
        # {<k1v1>, <k2v2>}
        ret = ['{']
        ind = 0
        for k, v in inp.items():
            # ''key': <value>'
            ret.append(f'{Color.key(wrapQuote(k))}: ')
            ret.extend(oneLineStringParts(v))
            if ind != len(inp) - 1:
                ret.append(', ')
            ind += 1
        ret.append('}')
        return ret
    elif type(inp) == list:
        # [i1, i2]
        ret = ['[']
        for ind, item in enumerate(inp):
            ret.extend(oneLineStringParts(item))
            if ind != len(inp) - 1:
                ret.append(', ')
        ret.append(']')
        return ret
    elif type(inp) == str:
        return [f'{Color.strval(wrapQuote(inp))}']
    elif inp == True:
        return ['true']
    elif inp == False:
        return ['false']
    elif type(inp) in {int, float}:
        return str(inp)
    elif inp is None:
        return [Color.null()]
    else:
        print(f'oneLineStringParts: Not sure know how to handle {inp}')
        sys.exit(1)


# Prints the given input in one line along with the pad string and the tail string.
def printOneLine(inp, pad, tail):
    dctstr = ''.join(oneLineStringParts(inp))
    print(f'{pad}{dctstr}{tail}')


def ifttt(cond, r1, r2):
    return r1 if cond else r2


# Pretty print the given object.
def pprint(inp, pad=None, tail=None):
    tail = tail or ''
    pad = pad or ''
    print_in_one_line_check = ifttt(
        isDctOrLst(inp),
        lenDct(inp) + len(pad) + len(tail) <= MAX_WIDTH,
        True)
    if print_in_one_line_check:
        printOneLine(inp, pad, tail)
        return
    if type(inp) == dict:
        print(pad + '{')
        cur_ind = 0
        for k, v in inp.items():
            key_and_baggage = f'{pad}  {Color.key(wrapQuote(k))}: '
            child_tail = ifttt(cur_ind != len(inp) - 1, ',', '')
            print_key_in_one_line_check = ifttt(
                isDctOrLst(v),
                len(key_and_baggage) + lenDct(v) + len(child_tail) < MAX_WIDTH,
                True)
            if print_key_in_one_line_check:
                printOneLine(v, pad=key_and_baggage, tail=child_tail)
            else:
                print(key_and_baggage)
                pprint(v, pad=pad + '  ', tail=child_tail)
            cur_ind += 1
        print(pad + '}' + tail)
    elif type(inp) == list:
        print(f'{pad}[')
        for cur_ind, item in enumerate(inp):
            child_tail = ifttt(cur_ind != len(inp) - 1, ',', '')
            pprint(item, pad=pad + '  ', tail=child_tail)
        print(f'{pad}]{tail}')
    else:
        print(f'pprint: Not sure what to do with {inp}')
        sys.exit(1)


def main():
    global ARGS
    global MAX_WIDTH
    parser = argparse.ArgumentParser()
    parser.add_argument('--nocache', action='store_true')
    parser.add_argument('--width', type=int, default=100)
    parser.add_argument('--debug', action='store_true')
    parser.add_argument('--nocolor', action='store_true')
    ARGS = parser.parse_args()
    MAX_WIDTH = ARGS.width
    pprint(json.load(sys.stdin))
    if ARGS.debug:
        print(f'hits:{LENGTH_CACHE_HITS} misses:{LENGTH_CACHE_MISSES}')


if __name__ == '__main__':
    main()