File: cbe.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 # cbe [options...] [filepaths/URIs...]
  27 #
  28 # Color By (file) Extension makes all lines look different, using different
  29 # ANSI-styles as it finds new trailing file-extensions. Since its list of
  30 # styles is limited, these will start being reused, given enough unique
  31 # extensions.
  32 #
  33 # Any ANSI-styles from the original input lines are ignored.
  34 #
  35 # The help option is `-h`, `--h`, `-help`, or `--help`.
  36 
  37 
  38 from re import compile
  39 from sys import argv, exit, stderr, stdin, stdout
  40 from typing import Dict, Tuple
  41 from urllib.request import urlopen
  42 
  43 
  44 # info is the message shown when the script isn't given any argument, or
  45 # when the leading argument is one of the standard cmd-line help options
  46 info = '''
  47 cbe [options...] [filepaths/URIs...]
  48 
  49 Color By (file) Extension makes all lines look different, using different
  50 ANSI-styles as it finds new trailing file-extensions. Since its list of
  51 styles is limited, these will start being reused, given enough unique
  52 extensions.
  53 
  54 Any ANSI-styles from the original input lines are ignored.
  55 
  56 The help option is `-h`, `--h`, `-help`, or `--help`.
  57 '''.strip()
  58 
  59 # a leading help-option arg means show the help message and quit
  60 if len(argv) == 2 and argv[1].lower() in ('-h', '--h', '-help', '--help'):
  61     print(info, file=stderr)
  62     exit(0)
  63 
  64 
  65 # ansi_re helps func style_line detect ANSI-style sequences
  66 ansi_re = compile('\x1b\\[[^m]+m')
  67 
  68 
  69 def style_line(w, s: str, styles: Tuple[str], ext2index: Dict[str, int]) -> None:
  70     '''Does what it says, using trailing file-extensions from lines.'''
  71 
  72     # ignore any ANSI-style sequences in the line
  73     s = ansi_re.sub('', s)
  74 
  75     # ignore trailing carriage-returns and line-feeds in lines, then
  76     # ignore trailing whitespace on what's left
  77     s = s.rstrip('\r\n').rstrip('\n').rstrip()
  78 
  79     ext = ''
  80     i = s.rfind('.')
  81     if i >= 0 and not s.endswith('.'):
  82         ext = s[i:]
  83     if ' ' in ext:
  84         ext = ''
  85 
  86     if ext != '':
  87         if ext in ext2index:
  88             w.write(styles[ext2index[ext]])
  89         else:
  90             ext2index[ext] = len(ext2index) % len(styles)
  91             w.write(styles[ext2index[ext]])
  92 
  93     w.write(s)
  94     w.write('\x1b[0m\n')
  95     w.flush()
  96 
  97 
  98 def seems_url(s: str) -> bool:
  99     for prot in ('https://', 'http://', 'file://', 'ftp://', 'data:'):
 100         if s.startswith(prot):
 101             return True
 102     return False
 103 
 104 
 105 # palette is the whole list of ANSI-styles used to make extensions stand out
 106 palette = [
 107     '\x1b[38;5;26m', # blue
 108     '\x1b[38;5;166m', # orange
 109     '\x1b[38;5;99m', # purple
 110     '\x1b[38;5;38m', # cyan
 111     '\x1b[38;5;213m', # pink
 112     '\x1b[38;5;29m', # green
 113     '\x1b[31m', # red
 114     '\x1b[38;5;249m', # gray
 115 
 116     # '\x1b[1m', # bold
 117     # '\x1b[4m', # underline
 118     # '\x1b[7m', # inverse
 119 ]
 120 
 121 ext2index = dict()
 122 
 123 try:
 124     args = argv[1:]
 125 
 126     if args.count('-') > 1:
 127         msg = 'reading from `-` (standard input) more than once not allowed'
 128         raise ValueError(msg)
 129 
 130     stdout.reconfigure(newline='\n', encoding='utf-8')
 131 
 132     # handle all named inputs given
 133     for path in args:
 134         if path == '-':
 135             for line in stdin:
 136                 style_line(stdout, line, palette, ext2index)
 137             continue
 138 
 139         if seems_url(path):
 140             with urlopen(path) as inp:
 141                 for line in inp:
 142                     line = str(line, encoding='utf-8')
 143                     style_line(stdout, line, palette, ext2index)
 144             continue
 145 
 146         with open(path) as inp:
 147             for line in inp:
 148                 style_line(stdout, line, palette, ext2index)
 149 
 150     # when no filenames are given, handle lines from stdin
 151     if len(args) == 0:
 152         for line in stdin:
 153             style_line(stdout, line, palette, ext2index)
 154 except BrokenPipeError:
 155     # quit quietly, instead of showing a confusing error message
 156     stderr.flush()
 157     stderr.close()
 158 except KeyboardInterrupt:
 159     # quit quietly, instead of showing a confusing error message
 160     stderr.flush()
 161     stderr.close()
 162     exit(2)
 163 except Exception as e:
 164     print(f'\x1b[31m{e}\x1b[0m', file=stderr)
 165     exit(1)