File: nt.py
   1 #!/usr/bin/python3
   2 
   3 # The MIT License (MIT)
   4 #
   5 # Copyright © 2020-2025 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 itertools import islice
  27 from math import isinf, isnan, modf, nan
  28 from sys import argv, exit, stderr, stdin, stdout
  29 from typing import Dict, Tuple
  30 
  31 
  32 info = '''
  33 nt [filepaths/URIs...]
  34 
  35 
  36 Nice Tables realigns and styles TSV (tab-separated values) data using ANSI
  37 sequences. When not given filepaths/URIs to read data from, this tool reads
  38 from standard input.
  39 '''
  40 
  41 # handle standard help cmd-line options, quitting right away in that case
  42 if len(argv) == 2 and argv[1] in ('-h', '--h', '-help', '--help'):
  43     print(info.strip())
  44     exit(0)
  45 
  46 
  47 # spaces is used by func emit_spaces
  48 spaces = tuple(' ' * i for i in range(64))
  49 
  50 
  51 def emit_spaces(w, n: int) -> None:
  52     'Speed-up writing spaces, by minimizing calls to func write.'
  53 
  54     if n < 1:
  55         return
  56 
  57     while n >= len(spaces):
  58         w.write(spaces[-1])
  59         n -= len(spaces)
  60     w.write(spaces[n])
  61 
  62 
  63 def count_decimals(s: str) -> int:
  64     '''
  65     Count the trailing decimal digits in a string, excluding the decimal dot.
  66     This func assumes the string can parse as a valid decimal-format float.
  67     '''
  68     k = s.find('.')
  69     return len(s) - k - 1 if k >= 0 else 0
  70 
  71 
  72 def get_table_info(src) -> Dict:
  73     'Gather metadata by reading TSV rows from the input given.'
  74 
  75     rows = []
  76     totals = []
  77     numeric = []
  78     max_decs = []
  79     max_widths = []
  80 
  81     def handle_row(items: Tuple[str]) -> None:
  82         nonlocal totals, numeric, max_decs, max_widths
  83 
  84         if len(max_widths) < len(items):
  85             extras = len(items) - len(max_widths)
  86             totals.extend([0.0] * extras)
  87             numeric.extend([False] * extras)
  88             max_decs.extend([0] * extras)
  89             max_widths.extend([0] * extras)
  90 
  91         for i, s in enumerate(items):
  92             max_widths[i] = max(max_widths[i], len(s))
  93             try:
  94                 f = float(s)
  95             except Exception:
  96                 continue
  97 
  98             if not (isnan(f) or isinf(f)):
  99                 max_decs[i] = max(max_decs[i], count_decimals(s))
 100                 totals[i] += f
 101                 numeric[i] = True
 102 
 103     for line in src:
 104         line = line.rstrip('\r\n').rstrip('\n')
 105         if line == '':
 106             continue
 107 
 108         items = tuple(line.split('\t'))
 109         rows.append(items)
 110         handle_row(items)
 111 
 112     # format column-totals as strings, and update column widths with those
 113     for i, num in enumerate(numeric):
 114         s = f'{totals[i]:.{max_decs[i]}f}' if num else '-'
 115         # s = f'{totals[i]:.{max_decs[i]}f}' if num else ''
 116         max_widths[i] = max(max_widths[i], len(s))
 117         totals[i] = s
 118 
 119     return {
 120         'rows': rows,
 121         'max-widths': tuple(max_widths),
 122         'max-decimals': tuple(max_decs),
 123         'totals': tuple(totals),
 124     }
 125 
 126 
 127 def show_table(w, meta: Dict) -> None:
 128     'Show/render table using the metadata given.'
 129 
 130     gap = len('  ')
 131     rows = meta['rows']
 132     max_widths = meta['max-widths']
 133     max_decs = meta['max-decimals']
 134     totals = meta['totals']
 135 
 136     if len(rows) == 0:
 137         return
 138 
 139     for row in rows:
 140         show_tiles(w, row, len(max_widths))
 141         show_row(w, row, max_widths, max_decs, gap)
 142 
 143     # make the totals line stand out by padding it with spaces, instead of
 144     # showing tiles: their lack should give a hint this final `row` isn't
 145     # part of the input
 146     emit_spaces(w, len(totals))
 147     show_row(w, totals, max_widths, max_decs, gap)
 148 
 149 
 150 def show_row(w, row: Tuple[str], widths: Tuple[int], max_decs: Tuple[int],
 151                 gap: int) -> None:
 152     pad = 0
 153 
 154     for i, s in enumerate(row):
 155         try:
 156             f = float(s)
 157         except Exception:
 158             f = nan
 159 
 160         if not (isnan(f) or isinf(f)):
 161             decs = count_decimals(s)
 162             if decs > 0:
 163                 trail = max_decs[i] - decs
 164             elif max_decs[i] > 0:
 165                 trail = max_decs[i] + 1
 166             else:
 167                 trail = 0
 168 
 169             lead = widths[i] - len(s)
 170             emit_spaces(w, max(pad + lead - trail, 0) + gap)
 171             if f > 0:
 172                 w.write(f'\x1b[38;5;29m{s}\x1b[0m')
 173             elif f < 0:
 174                 w.write(f'\x1b[38;5;1m{s}\x1b[0m')
 175             elif f == 0:
 176                 # w.write(f'\x1b[38;5;33m{s}\x1b[0m')
 177                 w.write(f'\x1b[38;5;26m{s}\x1b[0m')
 178             else:
 179                 w.write(s)
 180             pad = trail
 181         else:
 182             emit_spaces(w, pad + gap)
 183             w.write(s)
 184             pad = widths[i] - len(s)
 185 
 186     # don't forget to end the row/line
 187     w.write('\n')
 188 
 189 
 190 def show_tiles(w, row: Tuple[str], expected: int) -> None:
 191     w.write('\x1b[48;5;255m')
 192     for s in row:
 193         try:
 194             f = float(s)
 195             rem, _ = modf(f)
 196             if f > 0:
 197                 w.write('\x1b[38;5;29m■' if rem != 0 else '\x1b[38;5;22m■')
 198             elif f < 0:
 199                 w.write('\x1b[38;5;1m■' if rem != 0 else '\x1b[38;5;167m■')
 200             elif f == 0:
 201                 # w.write('\x1b[38;5;33m■')
 202                 w.write('\x1b[38;5;26m■')
 203             else:
 204                 w.write('\x1b[0m■')
 205         except Exception:
 206             if s == '':
 207                 w.write('\x1b[0m\x1b[48;5;255m○')
 208             else:
 209                 w.write('\x1b[38;5;244m■')
 210 
 211     if len(row) - expected:
 212         w.write('\x1b[0m\x1b[48;5;255m')
 213     for _ in range(expected - len(row)):
 214         w.write('×')
 215     w.write('\x1b[0m')
 216 
 217 
 218 def handle_tsv(w, src) -> None:
 219     show_table(w, get_table_info(src))
 220 
 221 
 222 def seems_url(s: str) -> bool:
 223     protocols = ('https://', 'http://', 'file://', 'ftp://', 'data:')
 224     return any(s.startswith(p) for p in protocols)
 225 
 226 
 227 try:
 228     if argv.count('-') > 1:
 229         msg = 'reading from `-` (standard input) more than once not allowed'
 230         raise ValueError(msg)
 231 
 232     if any(seems_url(e) for e in argv):
 233         from io import TextIOWrapper
 234         from urllib.request import urlopen
 235 
 236     for path in islice(argv, 1, None):
 237         if path == '-':
 238             handle_tsv(stdout, stdin)
 239         elif seems_url(path):
 240             with urlopen(path) as inp:
 241                 with TextIOWrapper(inp, encoding='utf-8') as txt:
 242                     handle_tsv(stdout, txt)
 243         else:
 244             with open(path, encoding='utf-8') as inp:
 245                 handle_tsv(stdout, inp)
 246 
 247     if len(argv) < 2:
 248         handle_tsv(stdout, stdin)
 249 except BrokenPipeError:
 250     # quit quietly, instead of showing a confusing error message
 251     stderr.close()
 252 except KeyboardInterrupt:
 253     exit(2)
 254 except Exception as e:
 255     print(f'\x1b[31m{e}\x1b[0m', file=stderr)
 256     exit(1)