summary refs log tree commit diff
path: root/dot_local/bin/jflatten.py
blob: 99849dcb6eb940880a9e4f7c9b0ef91b49b5204e (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
#!/usr/bin/env python3

import argparse
import decimal
import json
import sys

GREEN = "\033[1;32m"
GREY = "\033[0;37m"
CLEAR = "\033[0m"
RED = "\033[1;31m"
YELLOW = "\033[1;33m"
BLUE = "\033[1;34m"
MAGENTA = "\033[1;35m"


def flatten(*path, obj):
    if isinstance(obj, list):
        yield path, obj
        for n, value in enumerate(obj):
            yield from flatten(*path, n, obj=value)
    elif isinstance(obj, dict):
        yield path, obj
        for key, value in obj.items():
            yield from flatten(*path, key, obj=value)
    else:
        yield path, obj


def fmt_path(*path, colour):
    parts = []
    if not colour:
        for part in path:
            parts.append(f"[{part}]" if isinstance(part, int) else f".{part}")
    else:
        for part in path:
            if isinstance(part, int):
                parts.append(f"{GREY}[{BLUE}{part}{GREY}]{CLEAR}")
            else:
                parts.append(f"{GREY}.{MAGENTA}{part}{CLEAR}")
    return "".join(parts)


def fmt_value(value, colour):
    if not colour:
        return {
            None: "null",
            True: "true",
            False: "false",
        }.get(value, str(value))
    if value is None:
        return f"{GREY}null{CLEAR}"
    if value is True:
        return f"{GREEN}true{CLEAR}"
    if value is False:
        return f"{RED}false{CLEAR}"
    if isinstance(value, str):
        return f"{YELLOW}{value}{CLEAR}"
    if isinstance(value, (int, float, decimal.Decimal)):
        return f"{BLUE}{value}{CLEAR}"
    return str(value)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--sort-keys", "-s", action="store_true")
    parser.add_argument("--json", "-j", action="store_true")
    parser.add_argument("--show-empty-collections", action="store_true")
    parser.add_argument("--colour", "-c", action="store_true")
    parser.add_argument("--no-colour", "-C", action="store_true")
    parser.add_argument("file", nargs="?")
    args = parser.parse_args()

    if args.file:
        with open(args.file) as f:
            data = json.load(f)
    else:
        data = json.load(sys.stdin)

    flattened = (
        (path, value)
        for path, value in flatten(obj=data)
        if not isinstance(value, (dict, list))
        or not value
        and args.show_empty_collections
    )

    if args.json:
        print(
            json.dumps(
                {fmt_path(*path, colour=False): value for path, value in flattened},
                indent=True,
                sort_keys=args.sort_keys,
            ),
        )
        return

    if args.sort_keys:
        flattened = sorted(flattened)
    colour = not args.no_colour and (args.colour or sys.stdout.isatty())
    flattened_fmt = (
        (fmt_path(*path, colour=colour), fmt_value(value, colour=colour))
        for path, value in flattened
    )

    if args.json:
        print(json.dumps(dict(flattened_fmt), indent=True))
    else:
        print("\n".join(f"{path}:{value}" for path, value in flattened_fmt))


if __name__ == "__main__":
    main()