pipe regex highlighter - highlights regex patterns in pipe with configurable colors

hilite.py

#!/usr/bin/env python3

# it's like `.. | grep 'pattern\|'`, but with more colors !
# it can be useful when viewing log files when piped with `less -R`

import sys
import re

# colors :: [(attribute, foreground, background)]
# list of color scheme for pattern, matched  against n'th regex.
# order within tuple does not really matters.
# you can add more color scheme as you want.
# see https://bluesock.org/~willkg/dev/ansi.html for simple code lookup.
# see http://www.lihaoyi.com/post/BuildyourownCommandLinewithANSIescapecodes.html for explanation on ansi color codes.
colors = [
  (1, 31),
  (1, 32),
  (1, 33),
]

# ansi-colorize given string.
# wrapping given string with escapte codes
def colorize(s, n):
  ansi = ';'.join(map(str, colors[n%len(colors)]))
  return '\x1b[%sm%s\x1b[0m' % (ansi, s)

# command line arguments are all regex pattern.
# any no. of pattern can be specified.
# the order matters, s.t. later patterns are matched only if the pattern
# does not straddle previous pattern matches.
# TODO use grep-regex or perl-regex instead of python-regex?
# FIXME error no handled when re.compile fails
# FIXME if any of given pattern matches the ansi escape sequence,
# e.g. 'm', it behaves unexpectedly.
# .. but still handful for simple task usage..
pats = list(map(re.compile, sys.argv[1:]))

#for line in input.readline():
for line in sys.stdin.readlines():
  for idx, pat in enumerate(pats):
    line = re.sub(pat, lambda m: colorize(m.group(0), idx), line)
  print(line, end='') # new line included in [line]