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 PREFIX_RE = re.compile('^<(.*?)>\s*(.*)', re.S)
  45
  46 INTERACTIVE = True
  47
  48 LOG_LEVEL = 1
  49 LOG_LEVELS = {
  50                 'quiet': 0,
  51                 'normal': 1,
  52                 'verbose': 2,
  53                 'debug': 3,
  54                 'never': 4,
  55              }
  56
  57 LOG_FILE_FP = None
  58 LOG_CONTENT = ''
  59 CATCHERR_BUFFILE_FD = -1
  60 CATCHERR_BUFFILE_PATH = None
  61 CATCHERR_SAVED_2 = -1
  62
  63 def _general_print(head, color, msg = None, stream = None, level = 'normal'):
  64     global LOG_CONTENT
  65     if not stream:
  66         stream = sys.stdout
  67
  68     if LOG_LEVELS[level] > LOG_LEVEL:
  69         # skip
  70         return
  71
  72     # encode raw 'unicode' str to utf8 encoded str
  73     if msg and isinstance(msg, unicode):
  74         msg = msg.encode('utf-8', 'ignore')
  75
  76     errormsg = ''
  77     if CATCHERR_BUFFILE_FD > 0:
  78         size = os.lseek(CATCHERR_BUFFILE_FD , 0, os.SEEK_END)
  79         os.lseek(CATCHERR_BUFFILE_FD, 0, os.SEEK_SET)
  80         errormsg = os.read(CATCHERR_BUFFILE_FD, size)
  81         os.ftruncate(CATCHERR_BUFFILE_FD, 0)
  82
  83     # append error msg to LOG
  84     if errormsg:
  85         LOG_CONTENT += errormsg
  86
  87     # append normal msg to LOG
  88     save_msg = msg.strip() if msg else None
  89     if save_msg:
  90         timestr = time.strftime("[%m/%d %H:%M:%S] ", time.localtime())
  91         LOG_CONTENT += timestr + save_msg + '\n'
  92
  93     if errormsg:
  94         _color_print('', NO_COLOR, errormsg, stream, level)
  95
  96     _color_print(head, color, msg, stream, level)
  97
  98 def _color_print(head, color, msg, stream, level):
  99     colored = True
 100     if color == NO_COLOR or \
 101        not stream.isatty() or \
 102        os.getenv('ANSI_COLORS_DISABLED') is not None:
 103         colored = False
 104
 105     if head.startswith('\r'):
 106         # need not \n at last
 107         newline = False
 108     else:
 109         newline = True
 110
 111     if colored:
 112         head = '\033[%dm%s:\033[0m ' %(color, head)
 113         if not newline:
 114             # ESC cmd to clear line
 115             head = '\033[2K' + head
 116     else:
 117         if head:
 118             head += ': '
 119             if head.startswith('\r'):
 120                 head = head.lstrip()
 121                 newline = True
 122
 123     if msg is not None:
 124         if isinstance(msg, unicode):
 125             msg = msg.encode('utf8', 'ignore')
 126
 127         stream.write('%s%s' % (head, msg))
 128         if newline:
 129             stream.write('\n')
 130
 131     stream.flush()
 132
 133 def _color_perror(head, color, msg, level = 'normal'):
 134     if CATCHERR_BUFFILE_FD > 0:
 135         _general_print(head, color, msg, sys.stdout, level)
 136     else:
 137         _general_print(head, color, msg, sys.stderr, level)
 138
 139 def _split_msg(head, msg):
 140     if isinstance(msg, list):
 141         msg = '\n'.join(map(str, msg))
 142
 143     if msg.startswith('\n'):
 144         # means print \n at first
 145         msg = msg.lstrip()
 146         head = '\n' + head
 147
 148     elif msg.startswith('\r'):
 149         # means print \r at first
 150         msg = msg.lstrip()
 151         head = '\r' + head
 152
 153     m = PREFIX_RE.match(msg)
 154     if m:
 155         head += ' <%s>' % m.group(1)
 156         msg = m.group(2)
 157
 158     return head, msg
 159
 160 def get_loglevel():
 161     return (k for k,v in LOG_LEVELS.items() if v==LOG_LEVEL).next()
 162
 163 def set_loglevel(level):
 164     global LOG_LEVEL
 165     if level not in LOG_LEVELS:
 166         # no effect
 167         return
 168
 169     LOG_LEVEL = LOG_LEVELS[level]
 170
 171 def set_interactive(mode=True):
 172     global INTERACTIVE
 173     if mode:
 174         INTERACTIVE = True
 175     else:
 176         INTERACTIVE = False
 177
 178 def log(msg=''):
 179     # log msg to LOG_CONTENT then save to logfile
 180     global LOG_CONTENT
 181     if msg:
 182         LOG_CONTENT += msg
 183
 184 def raw(msg=''):
 185     _general_print('', NO_COLOR, msg)
 186
 187 def info(msg):
 188     head, msg = _split_msg('Info', msg)
 189     _general_print(head, INFO_COLOR, msg)
 190
 191 def verbose(msg):
 192     head, msg = _split_msg('Verbose', msg)
 193     _general_print(head, INFO_COLOR, msg, level = 'verbose')
 194
 195 def warning(msg):
 196     head, msg = _split_msg('Warning', msg)
 197     _color_perror(head, WARN_COLOR, msg)
 198
 199 def debug(msg):
 200     head, msg = _split_msg('Debug', msg)
 201     _color_perror(head, ERR_COLOR, msg, level = 'debug')
 202
 203 def error(msg):
 204     head, msg = _split_msg('Error', msg)
 205     _color_perror(head, ERR_COLOR, msg)
 206     sys.exit(1)
 207
 208 def ask(msg, default=True):
 209     _general_print('\rQ', ASK_COLOR, '')
 210     try:
 211         if default:
 212             msg += '(Y/n) '
 213         else:
 214             msg += '(y/N) '
 215         if INTERACTIVE:
 216             while True:
 217                 repl = raw_input(msg)
 218                 if repl.lower() == 'y':
 219                     return True
 220                 elif repl.lower() == 'n':
 221                     return False
 222                 elif not repl.strip():
 223                     # <Enter>
 224                     return default
 225
 226                 # else loop
 227         else:
 228             if default:
 229                 msg += ' Y'
 230             else:
 231                 msg += ' N'
 232             _general_print('', NO_COLOR, msg)
 233
 234             return default
 235     except KeyboardInterrupt:
 236         sys.stdout.write('\n')
 237         sys.exit(2)
 238
 239 def choice(msg, choices, default=0):
 240     if default >= len(choices):
 241         return None
 242     _general_print('\rQ', ASK_COLOR, '')
 243     try:
 244         msg += " [%s] " % '/'.join(choices)
 245         if INTERACTIVE:
 246             while True:
 247                 repl = raw_input(msg)
 248                 if repl in choices:
 249                     return repl
 250                 elif not repl.strip():
 251                     return choices[default]
 252         else:
 253             msg += choices[default]
 254             _general_print('', NO_COLOR, msg)
 255
 256             return choices[default]
 257     except KeyboardInterrupt:
 258         sys.stdout.write('\n')
 259         sys.exit(2)
 260
 261 def pause(msg=None):
 262     if INTERACTIVE:
 263         _general_print('\rQ', ASK_COLOR, '')
 264         if msg is None:
 265             msg = 'press <ENTER> to continue ...'
 266         raw_input(msg)
 267
 268 def set_logfile(fpath):
 269     global LOG_FILE_FP
 270
 271     def _savelogf():
 272         if 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 = 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