mic/msger.py

   1 #!/usr/bin/python -tt
   2 # vim: ai ts=4 sts=4 et sw=4
   3 #
   4 # Copyright (c) 2009, 2010, 2011 Intel, Inc.
   5 #
   6 # This program is free software; you can redistribute it and/or modify it
   7 # under the terms of the GNU General Public License as published by the Free
   8 # Software Foundation; version 2 of the License
   9 #
  10 # This program is distributed in the hope that it will be useful, but
  11 # WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
  12 # or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
  13 # for more details.
  14 #
  15 # You should have received a copy of the GNU General Public License along
  16 # with this program; if not, write to the Free Software Foundation, Inc., 59
  17 # Temple Place - Suite 330, Boston, MA 02111-1307, USA.
  18
  19 import os,sys
  20 import re
  21 import time
  22
  23 __ALL__ = ['set_mode',
  24            'get_loglevel',
  25            'set_loglevel',
  26            'set_logfile',
  27            'raw',
  28            'debug',
  29            'verbose',
  30            'info',
  31            'warning',
  32            'error',
  33            'ask',
  34            'pause',
  35           ]
  36
  37 # COLORs in ANSI
  38 INFO_COLOR = 32 # green
  39 WARN_COLOR = 33 # yellow
  40 ERR_COLOR  = 31 # red
  41 ASK_COLOR  = 34 # blue
  42 NO_COLOR = 0
  43
  44 HOST_TIMEZONE = time.timezone
  45
  46 PREFIX_RE = re.compile('^<(.*?)>\s*(.*)', re.S)
  47
  48 INTERACTIVE = True
  49
  50 LOG_LEVEL = 1
  51 LOG_LEVELS = {
  52                 'quiet': 0,
  53                 'normal': 1,
  54                 'verbose': 2,
  55                 'debug': 3,
  56                 'never': 4,
  57              }
  58
  59 LOG_FILE_FP = None
  60 LOG_CONTENT = ''
  61 CATCHERR_BUFFILE_FD = -1
  62 CATCHERR_BUFFILE_PATH = None
  63 CATCHERR_SAVED_2 = -1
  64
  65 def _general_print(head, color, msg = None, stream = None, level = 'normal'):
  66     global LOG_CONTENT
  67     if not stream:
  68         stream = sys.stdout
  69
  70     if LOG_LEVELS[level] > LOG_LEVEL:
  71         # skip
  72         return
  73
  74     # encode raw 'unicode' str to utf8 encoded str
  75     if msg and isinstance(msg, unicode):
  76         msg = msg.encode('utf-8', 'ignore')
  77
  78     errormsg = ''
  79     if CATCHERR_BUFFILE_FD > 0:
  80         size = os.lseek(CATCHERR_BUFFILE_FD , 0, os.SEEK_END)
  81         os.lseek(CATCHERR_BUFFILE_FD, 0, os.SEEK_SET)
  82         errormsg = os.read(CATCHERR_BUFFILE_FD, size)
  83         os.ftruncate(CATCHERR_BUFFILE_FD, 0)
  84
  85     # append error msg to LOG
  86     if errormsg:
  87         LOG_CONTENT += errormsg
  88
  89     # append normal msg to LOG
  90     save_msg = msg.strip() if msg else None
  91     if save_msg:
  92         global HOST_TIMEZONE
  93         timestr = time.strftime("[%m/%d %H:%M:%S] ",
  94                                 time.gmtime(time.time() - HOST_TIMEZONE))
  95         LOG_CONTENT += timestr + save_msg + '\n'
  96
  97     if errormsg:
  98         _color_print('', NO_COLOR, errormsg, stream, level)
  99
 100     _color_print(head, color, msg, stream, level)
 101
 102 def _color_print(head, color, msg, stream, level):
 103     colored = True
 104     if color == NO_COLOR or \
 105        not stream.isatty() or \
 106        os.getenv('ANSI_COLORS_DISABLED') is not None:
 107         colored = False
 108
 109     if head.startswith('\r'):
 110         # need not \n at last
 111         newline = False
 112     else:
 113         newline = True
 114
 115     if colored:
 116         head = '\033[%dm%s:\033[0m ' %(color, head)
 117         if not newline:
 118             # ESC cmd to clear line
 119             head = '\033[2K' + head
 120     else:
 121         if head:
 122             head += ': '
 123             if head.startswith('\r'):
 124                 head = head.lstrip()
 125                 newline = True
 126
 127     if msg is not None:
 128         if isinstance(msg, unicode):
 129             msg = msg.encode('utf8', 'ignore')
 130
 131         stream.write('%s%s' % (head, msg))
 132         if newline:
 133             stream.write('\n')
 134
 135     stream.flush()
 136
 137 def _color_perror(head, color, msg, level = 'normal'):
 138     if CATCHERR_BUFFILE_FD > 0:
 139         _general_print(head, color, msg, sys.stdout, level)
 140     else:
 141         _general_print(head, color, msg, sys.stderr, level)
 142
 143 def _split_msg(head, msg):
 144     if isinstance(msg, list):
 145         msg = '\n'.join(map(str, msg))
 146
 147     if msg.startswith('\n'):
 148         # means print \n at first
 149         msg = msg.lstrip()
 150         head = '\n' + head
 151
 152     elif msg.startswith('\r'):
 153         # means print \r at first
 154         msg = msg.lstrip()
 155         head = '\r' + head
 156
 157     m = PREFIX_RE.match(msg)
 158     if m:
 159         head += ' <%s>' % m.group(1)
 160         msg = m.group(2)
 161
 162     return head, msg
 163
 164 def get_loglevel():
 165     return (k for k,v in LOG_LEVELS.items() if v==LOG_LEVEL).next()
 166
 167 def set_loglevel(level):
 168     global LOG_LEVEL
 169     if level not in LOG_LEVELS:
 170         # no effect
 171         return
 172
 173     LOG_LEVEL = LOG_LEVELS[level]
 174
 175 def set_interactive(mode=True):
 176     global INTERACTIVE
 177     if mode:
 178         INTERACTIVE = True
 179     else:
 180         INTERACTIVE = False
 181
 182 def raw(msg=''):
 183     _general_print('', NO_COLOR, msg)
 184
 185 def info(msg):
 186     head, msg = _split_msg('Info', msg)
 187     _general_print(head, INFO_COLOR, msg)
 188
 189 def verbose(msg):
 190     head, msg = _split_msg('Verbose', msg)
 191     _general_print(head, INFO_COLOR, msg, level = 'verbose')
 192
 193 def warning(msg):
 194     head, msg = _split_msg('Warning', msg)
 195     _color_perror(head, WARN_COLOR, msg)
 196
 197 def debug(msg):
 198     head, msg = _split_msg('Debug', msg)
 199     _color_perror(head, ERR_COLOR, msg, level = 'debug')
 200
 201 def error(msg):
 202     head, msg = _split_msg('Error', msg)
 203     _color_perror(head, ERR_COLOR, msg)
 204     sys.exit(1)
 205
 206 def ask(msg, default=True):
 207     _general_print('\rQ', ASK_COLOR, '')
 208     try:
 209         if default:
 210             msg += '(Y/n) '
 211         else:
 212             msg += '(y/N) '
 213         if INTERACTIVE:
 214             while True:
 215                 repl = raw_input(msg)
 216                 if repl.lower() == 'y':
 217                     return True
 218                 elif repl.lower() == 'n':
 219                     return False
 220                 elif not repl.strip():
 221                     # <Enter>
 222                     return default
 223
 224                 # else loop
 225         else:
 226             if default:
 227                 msg += ' Y'
 228             else:
 229                 msg += ' N'
 230             _general_print('', NO_COLOR, msg)
 231
 232             return default
 233     except KeyboardInterrupt:
 234         sys.stdout.write('\n')
 235         sys.exit(2)
 236
 237 def choice(msg, choices, default=0):
 238     if default >= len(choices):
 239         return None
 240     _general_print('\rQ', ASK_COLOR, '')
 241     try:
 242         msg += " [%s] " % '/'.join(choices)
 243         if INTERACTIVE:
 244             while True:
 245                 repl = raw_input(msg)
 246                 if repl in choices:
 247                     return repl
 248                 elif not repl.strip():
 249                     return choices[default]
 250         else:
 251             msg += choices[default]
 252             _general_print('', NO_COLOR, msg)
 253
 254             return choices[default]
 255     except KeyboardInterrupt:
 256         sys.stdout.write('\n')
 257         sys.exit(2)
 258
 259 def pause(msg=None):
 260     if INTERACTIVE:
 261         _general_print('\rQ', ASK_COLOR, '')
 262         if msg is None:
 263             msg = 'press <ENTER> to continue ...'
 264         raw_input(msg)
 265
 266 def set_logfile(fpath):
 267     global LOG_FILE_FP
 268
 269     def _savelogf():
 270         if LOG_FILE_FP:
 271             if not os.path.exists(os.path.dirname(LOG_FILE_FP)):
 272                 os.makedirs(os.path.dirname(LOG_FILE_FP))
 273             fp = open(LOG_FILE_FP, 'w')
 274             fp.write(LOG_CONTENT)
 275             fp.close()
 276
 277     if LOG_FILE_FP is not None:
 278         warning('duplicate log file configuration')
 279
 280     LOG_FILE_FP = os.path.abspath(os.path.expanduser(fpath))
 281
 282     import atexit
 283     atexit.register(_savelogf)
 284
 285 def enable_logstderr(fpath):
 286     global CATCHERR_BUFFILE_FD
 287     global CATCHERR_BUFFILE_PATH
 288     global CATCHERR_SAVED_2
 289
 290     if os.path.exists(fpath):
 291         os.remove(fpath)
 292     CATCHERR_BUFFILE_PATH = fpath
 293     CATCHERR_BUFFILE_FD = os.open(CATCHERR_BUFFILE_PATH, os.O_RDWR|os.O_CREAT)
 294     CATCHERR_SAVED_2 = os.dup(2)
 295     os.dup2(CATCHERR_BUFFILE_FD, 2)
 296
 297 def disable_logstderr():
 298     global CATCHERR_BUFFILE_FD
 299     global CATCHERR_BUFFILE_PATH
 300     global CATCHERR_SAVED_2
 301
 302     raw(msg = None) # flush message buffer and print it.
 303     os.dup2(CATCHERR_SAVED_2, 2)
 304     os.close(CATCHERR_SAVED_2)
 305     os.close(CATCHERR_BUFFILE_FD)
 306     os.unlink(CATCHERR_BUFFILE_PATH)
 307     CATCHERR_BUFFILE_FD = -1
 308     CATCHERR_BUFFILE_PATH = None
 309     CATCHERR_SAVED_2 = -1