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 log(msg=''):
 183     # log msg to LOG_CONTENT then save to logfile
 184     global LOG_CONTENT
 185     if msg:
 186         LOG_CONTENT += msg
 187
 188 def raw(msg=''):
 189     _general_print('', NO_COLOR, msg)
 190
 191 def info(msg):
 192     head, msg = _split_msg('Info', msg)
 193     _general_print(head, INFO_COLOR, msg)
 194
 195 def verbose(msg):
 196     head, msg = _split_msg('Verbose', msg)
 197     _general_print(head, INFO_COLOR, msg, level = 'verbose')
 198
 199 def warning(msg):
 200     head, msg = _split_msg('Warning', msg)
 201     _color_perror(head, WARN_COLOR, msg)
 202
 203 def debug(msg):
 204     head, msg = _split_msg('Debug', msg)
 205     _color_perror(head, ERR_COLOR, msg, level = 'debug')
 206
 207 def error(msg):
 208     head, msg = _split_msg('Error', msg)
 209     _color_perror(head, ERR_COLOR, msg)
 210     sys.exit(1)
 211
 212 def ask(msg, default=True):
 213     _general_print('\rQ', ASK_COLOR, '')
 214     try:
 215         if default:
 216             msg += '(Y/n) '
 217         else:
 218             msg += '(y/N) '
 219         if INTERACTIVE:
 220             while True:
 221                 repl = raw_input(msg)
 222                 if repl.lower() == 'y':
 223                     return True
 224                 elif repl.lower() == 'n':
 225                     return False
 226                 elif not repl.strip():
 227                     # <Enter>
 228                     return default
 229
 230                 # else loop
 231         else:
 232             if default:
 233                 msg += ' Y'
 234             else:
 235                 msg += ' N'
 236             _general_print('', NO_COLOR, msg)
 237
 238             return default
 239     except KeyboardInterrupt:
 240         sys.stdout.write('\n')
 241         sys.exit(2)
 242
 243 def choice(msg, choices, default=0):
 244     if default >= len(choices):
 245         return None
 246     _general_print('\rQ', ASK_COLOR, '')
 247     try:
 248         msg += " [%s] " % '/'.join(choices)
 249         if INTERACTIVE:
 250             while True:
 251                 repl = raw_input(msg)
 252                 if repl in choices:
 253                     return repl
 254                 elif not repl.strip():
 255                     return choices[default]
 256         else:
 257             msg += choices[default]
 258             _general_print('', NO_COLOR, msg)
 259
 260             return choices[default]
 261     except KeyboardInterrupt:
 262         sys.stdout.write('\n')
 263         sys.exit(2)
 264
 265 def pause(msg=None):
 266     if INTERACTIVE:
 267         _general_print('\rQ', ASK_COLOR, '')
 268         if msg is None:
 269             msg = 'press <ENTER> to continue ...'
 270         raw_input(msg)
 271
 272 def set_logfile(fpath):
 273     global LOG_FILE_FP
 274
 275     def _savelogf():
 276         if LOG_FILE_FP:
 277             fp = open(LOG_FILE_FP, 'w')
 278             fp.write(LOG_CONTENT)
 279             fp.close()
 280
 281     if LOG_FILE_FP is not None:
 282         warning('duplicate log file configuration')
 283
 284     LOG_FILE_FP = fpath
 285
 286     import atexit
 287     atexit.register(_savelogf)
 288
 289 def enable_logstderr(fpath):
 290     global CATCHERR_BUFFILE_FD
 291     global CATCHERR_BUFFILE_PATH
 292     global CATCHERR_SAVED_2
 293
 294     if os.path.exists(fpath):
 295         os.remove(fpath)
 296     CATCHERR_BUFFILE_PATH = fpath
 297     CATCHERR_BUFFILE_FD = os.open(CATCHERR_BUFFILE_PATH, os.O_RDWR|os.O_CREAT)
 298     CATCHERR_SAVED_2 = os.dup(2)
 299     os.dup2(CATCHERR_BUFFILE_FD, 2)
 300
 301 def disable_logstderr():
 302     global CATCHERR_BUFFILE_FD
 303     global CATCHERR_BUFFILE_PATH
 304     global CATCHERR_SAVED_2
 305
 306     raw(msg = None) # flush message buffer and print it.
 307     os.dup2(CATCHERR_SAVED_2, 2)
 308     os.close(CATCHERR_SAVED_2)
 309     os.close(CATCHERR_BUFFILE_FD)
 310     os.unlink(CATCHERR_BUFFILE_PATH)
 311     CATCHERR_BUFFILE_FD = -1
 312     CATCHERR_BUFFILE_PATH = None
 313     CATCHERR_SAVED_2 = -1