Home | History | Annotate | Download | only in common 
      1 # Copyright 2010 Google Inc. All Rights Reserved.
      2 
      3 from itertools import chain
      4 import gzip
      5 import logging
      6 import logging.handlers
      7 import time
      8 import traceback
      9 
     10 
     11 def SetUpRootLogger(filename=None, level=None, display_flags={}):
     12   console_handler = logging.StreamHandler()
     13   console_handler.setFormatter(CustomFormatter(AnsiColorCoder(), display_flags))
     14   logging.root.addHandler(console_handler)
     15 
     16   if filename:
     17     file_handler = logging.handlers.RotatingFileHandler(
     18         filename,
     19         maxBytes=10 * 1024 * 1024,
     20         backupCount=9,
     21         delay=True)
     22     file_handler.setFormatter(CustomFormatter(NullColorCoder(), display_flags))
     23     logging.root.addHandler(file_handler)
     24 
     25   if level:
     26     logging.root.setLevel(level)
     27 
     28 
     29 class NullColorCoder(object):
     30 
     31   def __call__(self, *args):
     32     return ''
     33 
     34 
     35 class AnsiColorCoder(object):
     36   CODES = {'reset': (0,),
     37            'bold': (1, 22),
     38            'italics': (3, 23),
     39            'underline': (4, 24),
     40            'inverse': (7, 27),
     41            'strikethrough': (9, 29),
     42            'black': (30, 40),
     43            'red': (31, 41),
     44            'green': (32, 42),
     45            'yellow': (33, 43),
     46            'blue': (34, 44),
     47            'magenta': (35, 45),
     48            'cyan': (36, 46),
     49            'white': (37, 47)}
     50 
     51   def __call__(self, *args):
     52     codes = []
     53 
     54     for arg in args:
     55       if arg.startswith('bg-') or arg.startswith('no-'):
     56         codes.append(self.CODES[arg[3:]][1])
     57       else:
     58         codes.append(self.CODES[arg][0])
     59 
     60     return '\033[%sm' % ';'.join(map(str, codes))
     61 
     62 
     63 class CustomFormatter(logging.Formatter):
     64   COLORS = {'DEBUG': ('white',),
     65             'INFO': ('green',),
     66             'WARN': ('yellow', 'bold'),
     67             'ERROR': ('red', 'bold'),
     68             'CRIT': ('red', 'inverse', 'bold')}
     69 
     70   def __init__(self, coder, display_flags={}):
     71     items = []
     72 
     73     if display_flags.get('datetime', True):
     74       items.append('%(asctime)s')
     75     if display_flags.get('level', True):
     76       items.append('%(levelname)s')
     77     if display_flags.get('name', True):
     78       items.append(coder('cyan') + '[%(threadName)s:%(name)s]' + coder('reset'))
     79     items.append('%(prefix)s%(message)s')
     80 
     81     logging.Formatter.__init__(self, fmt=' '.join(items))
     82 
     83     self._coder = coder
     84 
     85   def formatTime(self, record):
     86     ct = self.converter(record.created)
     87     t = time.strftime('%Y-%m-%d %H:%M:%S', ct)
     88     return '%s.%02d' % (t, record.msecs / 10)
     89 
     90   def formatLevelName(self, record):
     91     if record.levelname in ['WARNING', 'CRITICAL']:
     92       levelname = record.levelname[:4]
     93     else:
     94       levelname = record.levelname
     95 
     96     return ''.join([self._coder(*self.COLORS[levelname]), levelname,
     97                     self._coder('reset')])
     98 
     99   def formatMessagePrefix(self, record):
    100     try:
    101       return ' %s%s:%s ' % (self._coder('black', 'bold'), record.prefix,
    102                             self._coder('reset'))
    103     except AttributeError:
    104       return ''
    105 
    106   def format(self, record):
    107     if record.exc_info:
    108       if not record.exc_text:
    109         record.exc_text = self.formatException(record.exc_info)
    110     else:
    111       record.exc_text = ''
    112 
    113     fmt = record.__dict__.copy()
    114     fmt.update({'levelname': self.formatLevelName(record),
    115                 'asctime': self.formatTime(record),
    116                 'prefix': self.formatMessagePrefix(record)})
    117 
    118     s = []
    119 
    120     for line in chain(record.getMessage().splitlines(),
    121                       record.exc_text.splitlines()):
    122       fmt['message'] = line
    123 
    124       s.append(self._fmt % fmt)
    125 
    126     return '\n'.join(s)
    127 
    128 
    129 class CompressedFileHandler(logging.FileHandler):
    130 
    131   def _open(self):
    132     return gzip.open(self.baseFilename + '.gz', self.mode, 9)
    133 
    134 
    135 def HandleUncaughtExceptions(fun):
    136   """Catches all exceptions that would go outside decorated fun scope."""
    137 
    138   def _Interceptor(*args, **kwargs):
    139     try:
    140       return fun(*args, **kwargs)
    141     except StandardError:
    142       logging.exception('Uncaught exception:')
    143 
    144   return _Interceptor
    145