logger.py revision f2a3ef46f75d2196a93d3ed27f4d1fcf22b54fbe 
1# Copyright 2010 Google Inc. All Rights Reserved.
2
3from itertools import chain
4import gzip
5import logging
6import logging.handlers
7import time
8import traceback
9
10
11def 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
29class NullColorCoder(object):
30
31  def __call__(self, *args):
32    return ''
33
34
35class 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
63class 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
129class CompressedFileHandler(logging.FileHandler):
130
131  def _open(self):
132    return gzip.open(self.baseFilename + '.gz', self.mode, 9)
133
134
135def 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