Highlight stdin or file input

#!/usr/bin/env python

import sys
import re

def highlight_text(text,pat):
    def replacement_funk(matchobj):  return '\x1b[1;33m%s\x1b[0m'%matchobj.group(0)
    return re.sub(pat,replacement_funk,text)

if __name__ == '__main__':
    if len(sys.argv) == 2:
        input = sys.stdin
        pat = sys.argv[1]
    elif len(sys.argv) == 3:
        input = open(sys.argv[2])
        pat = sys.argv[1]
    else:
        sys.stderr.write("colorme pattern [inputfile]")
    text = input.read()
    print highlight_text(text,pat)