python - mini grep with text highlight

grep58.py

#! /usr/bin/python

"""
Usage:
grep58.py <pattern> <file>...
grep58.py <pattern>

Examples:

# from file(s)
./grep58.py "tench" ~/46/*
./grep58.py "tench|gudgeon" ~/logs/*.log
./grep58.py "exception|timeout" 1.log 2.log

# from directory
./grep58.py perch ~/log_dir/*

# from directories
./grep58.py barbel ~/log_dir/* ~/Music/*

# from <stdin>
df -h | ./grep58.py dev
"""

import sys, re, fileinput

BLUE = '\033[94m'
GREEN = '\033[92m'
YELLOW = '\033[93m'
RED = '\033[91m'
END = '\033[0m'
BOLD = "\033[1m"

def highlight_text(obj):
 text = obj.group(0)
 return BOLD + RED + text + END

def main():

 pattern = re.compile(sys.argv[1])

 for line in fileinput.input(sys.argv[2:]):
  if pattern.search(line):
   color_text = re.sub((pattern), highlight_text, line.rstrip())
   if fileinput.isstdin():
    fmt = '{lineno}:{line}'
   else:
    fmt = '{filename}:[{lineno}]:{line}'
  
   # print matched lines
   print fmt.format(filename=fileinput.filename(), lineno=fileinput.filelineno(), line=color_text)

if __name__ == "__main__":
 
 if len(sys.argv) < 2: # need a <pattern> at least
  print (__doc__)
  sys.exit(58)  

 main()