File: nj.py
   1 #!/usr/bin/python3
   2 
   3 # The MIT License (MIT)
   4 #
   5 # Copyright © 2024 pacman64
   6 #
   7 # Permission is hereby granted, free of charge, to any person obtaining a copy
   8 # of this software and associated documentation files (the “Software”), to deal
   9 # in the Software without restriction, including without limitation the rights
  10 # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
  11 # copies of the Software, and to permit persons to whom the Software is
  12 # furnished to do so, subject to the following conditions:
  13 #
  14 # The above copyright notice and this permission notice shall be included in
  15 # all copies or substantial portions of the Software.
  16 #
  17 # THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
  18 # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
  19 # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
  20 # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
  21 # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
  22 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
  23 # SOFTWARE.
  24 
  25 
  26 from json import dumps, load
  27 from math import modf
  28 from sys import argv, exit, stderr, stdin, stdout
  29 
  30 
  31 info = '''
  32 nj [filepath/URI...]
  33 
  34 
  35 Nice Json indents and styles JSON data. When not given a filepath to read
  36 data from, this tool reads from standard input.
  37 '''
  38 
  39 if len(argv) == 2 and argv[1] in ('-h', '--h', '-help', '--help'):
  40     print(info.strip())
  41     exit(0)
  42 
  43 
  44 # level2indent speeds up func indent
  45 level2indent = tuple(i * '  ' for i in range(64))
  46 
  47 
  48 def indent(w, level: int) -> None:
  49     'Emit 2 spaces for each indentation level.'
  50 
  51     if level < 1:
  52         return
  53 
  54     if level < len(level2indent):
  55         w.write(level2indent[level])
  56         return
  57 
  58     while level >= len(level2indent):
  59         w.write(level2indent[-1])
  60         level -= len(level2indent)
  61     w.write(level2indent[level])
  62 
  63 
  64 def restyle(w, data, pre: int = 0, level: int = 0) -> None:
  65     'This func is where all the recursive output-action starts/happens.'
  66 
  67     f = type2func[type(data)]
  68     if f != None:
  69         f(w, data, pre, level)
  70         # don't forget to end the final line, if at `top-level`
  71         if level == 0:
  72             w.write('\n')
  73     else:
  74         raise ValueError(f'unsupported type {type(data)}')
  75 
  76 
  77 def restyle_bool(w, data: bool, pre: int, level: int) -> None:
  78     indent(w, pre)
  79     if data:
  80         w.write('\x1b[38;5;74mtrue\x1b[0m')
  81     else:
  82         w.write('\x1b[38;5;74mfalse\x1b[0m')
  83 
  84 
  85 def restyle_dict(w, data: dict, pre: int, level: int) -> None:
  86     if len(data) == 0:
  87         indent(w, pre)
  88         w.write('\x1b[38;5;248m{}\x1b[0m')
  89         return
  90 
  91     indent(w, pre)
  92     w.write('\x1b[38;5;248m{\x1b[0m\n')
  93     for (i, k) in enumerate(data):
  94         if i > 0:
  95             w.write('\x1b[38;5;248m,\x1b[0m\n')
  96         restyle_key(w, k, level + 1)
  97         restyle(w, data[k], 0, level + 1)
  98     w.write('\n')
  99     indent(w, level)
 100     w.write('\x1b[38;5;248m}\x1b[0m')
 101 
 102 
 103 def pick_numeric_style(n) -> str:
 104     if n > 0:
 105         return '\x1b[38;5;29m'
 106     if n < 0:
 107         return '\x1b[38;5;1m'
 108     return '\x1b[38;5;26m'
 109 
 110 
 111 def restyle_float(w, data: float, pre: int, level: int) -> None:
 112     indent(w, pre)
 113     # can't figure out a direct float-format option which both avoids the
 114     # chance of scientific notation (invalid as JSON), and the chance of
 115     # unneeded trailing decimal 0s; neither {data:-1} nor {data:-1f} work
 116     if modf(data)[0] == 0.0:
 117         w.write(f'{pick_numeric_style(data)}{data}\x1b[0m')
 118     else:
 119         s = f'{data:-1f}'.rstrip('0')
 120         w.write(f'{pick_numeric_style(data)}{s}\x1b[0m')
 121 
 122 
 123 def restyle_int(w, data: int, pre: int, level: int) -> None:
 124     indent(w, pre)
 125     w.write(f'{pick_numeric_style(data)}{data}\x1b[0m')
 126 
 127 
 128 def restyle_key(w, k: str, level: int) -> None:
 129     indent(w, level)
 130     f = dumps
 131     w.write(f'\x1b[38;5;248m"\x1b[38;5;99m{f(k)[1:-1]}\x1b[38;5;248m": \x1b[0m')
 132 
 133 
 134 def restyle_list(w, data: list, pre: int, level: int) -> None:
 135     if len(data) == 0:
 136         indent(w, pre)
 137         w.write('\x1b[38;5;248m[]\x1b[0m')
 138         return
 139 
 140     indent(w, pre)
 141     w.write('\x1b[38;5;248m[\x1b[0m\n')
 142     for (i, e) in enumerate(data):
 143         if i > 0:
 144             w.write('\x1b[38;5;248m,\x1b[0m\n')
 145         restyle(w, e, level + 1, level + 1)
 146     w.write('\n')
 147     indent(w, level)
 148     w.write('\x1b[38;5;248m]\x1b[0m')
 149 
 150 
 151 def restyle_none(w, data: bool, pre: int, level: int) -> None:
 152     indent(w, pre)
 153     w.write('\x1b[38;5;248mnull\x1b[0m')
 154 
 155 
 156 def restyle_str(w, data: str, pre: int, level: int) -> None:
 157     indent(w, pre)
 158     # f'\x1b[38;5;248m"\x1b[38;5;24m{dumps(data)[1:-1]}\x1b[38;5;248m"\x1b[0m'
 159     w.write(f'\x1b[38;5;248m"\x1b[0m{dumps(data)[1:-1]}\x1b[38;5;248m"\x1b[0m')
 160 
 161 
 162 def seems_url(s: str) -> bool:
 163     protocols = ('https://', 'http://', 'file://', 'ftp://', 'data:')
 164     return any(s.startswith(p) for p in protocols)
 165 
 166 
 167 # type2func helps func restyle dispatch styler funcs, using values' types
 168 # as the keys; while JSON doesn't have int values, the stdlib auto-promotes
 169 # round JSON numbers into ints, and so this table has a func for those too
 170 type2func = {
 171     type(None): restyle_none,
 172     bool: restyle_bool,
 173     int: restyle_int,
 174     float: restyle_float,
 175     str: restyle_str,
 176     list: restyle_list,
 177     tuple: restyle_list,
 178     dict: restyle_dict,
 179 }
 180 
 181 
 182 try:
 183     if len(argv) < 2:
 184         restyle(stdout, load(stdin))
 185     elif len(argv) == 2:
 186         name = argv[1]
 187         if name == '-':
 188             restyle(stdout, load(stdin))
 189         elif seems_url(name):
 190             from urllib.request import urlopen
 191             with urlopen(name) as inp:
 192                 restyle(stdout, load(inp))
 193         else:
 194             with open(name, encoding='utf-8') as inp:
 195                 restyle(stdout, load(inp))
 196     else:
 197         raise ValueError('multiple inputs not allowed')
 198 except BrokenPipeError:
 199     # quit quietly, instead of showing a confusing error message
 200     stderr.close()
 201     exit(0)
 202 except KeyboardInterrupt:
 203     exit(2)
 204 except Exception as e:
 205     print(f'\x1b[31m{e}\x1b[0m', file=stderr)
 206     exit(1)