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