Add colourful formatting options to jflatten.py

1 files changed, 69 insertions, 6 deletions

diff --git a/dot_local/bin/jflatten.py b/dot_local/bin/jflatten.py
index bacc754..99849dc 100755
--- a/dot_local/bin/jflatten.py
+++ b/dot_local/bin/jflatten.py
@@ -1,31 +1,73 @@
 #!/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 path_str(*path):
-    return "".join(
-        f"[{part}]" if isinstance(part, int) else f".{part}" for part in path
-    )
+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()
 
@@ -35,10 +77,31 @@ def main():
     else:
         data = json.load(sys.stdin)
 
-    flattened = flatten(obj=data)
+    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)
-    flattened_fmt = ((path_str(*path), value) for path, value in 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))