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     if LOG_FILE_FP:
  86         if errormsg:
  87             LOG_CONTENT += errormsg
  88
  89         save_msg = msg.strip() if msg else None
  90         if save_msg:
  91             global HOST_TIMEZONE
  92             timestr = time.strftime("[%m/%d %H:%M:%S] ",
  93                                     time.gmtime(time.time() - HOST_TIMEZONE))
  94             LOG_CONTENT += timestr + save_msg + '\n'
  95
  96     if errormsg:
  97         _color_print('', NO_COLOR, errormsg, stream, level)
  98
  99     _color_print(head, color, msg, stream, level)
 100
 101 def _color_print(head, color, msg, stream, level):
 102     colored = True
 103     if color == NO_COLOR or \
 104        not stream.isatty() or \
 105        os.getenv('ANSI_COLORS_DISABLED') is not None:
 106         colored = False
 107
 108     if head.startswith('\r'):
 109         # need not \n at last
 110         newline = False
 111     else:
 112         newline = True
 113
 114     if colored:
 115         head = '\033[%dm%s:\033[0m ' %(color, head)
 116         if not newline:
 117             # ESC cmd to clear line
 118             head = '\033[2K' + head
 119     else:
 120         if head:
 121             head += ': '
 122             if head.startswith('\r'):
 123                 head = head.lstrip()
 124                 newline = True
 125
 126     if msg is not None:
 127         if isinstance(msg, unicode):
 128             msg = msg.encode('utf8', 'ignore')
 129
 130         stream.write('%s%s' % (head, msg))
 131         if newline:
 132             stream.write('\n')
 133
 134     stream.flush()
 135
 136 def _color_perror(head, color, msg, level = 'normal'):
 137     if CATCHERR_BUFFILE_FD > 0:
 138         _general_print(head, color, msg, sys.stdout, level)
 139     else:
 140         _general_print(head, color, msg, sys.stderr, level)
 141
 142 def _split_msg(head, msg):
 143     if isinstance(msg, list):
 144         msg = '\n'.join(map(str, msg))
 145
 146     if msg.startswith('\n'):
 147         # means print \n at first
 148         msg = msg.lstrip()
 149         head = '\n' + head
 150
 151     elif msg.startswith('\r'):
 152         # means print \r at first
 153         msg = msg.lstrip()
 154         head = '\r' + head
 155
 156     m = PREFIX_RE.match(msg)
 157     if m:
 158         head += ' <%s>' % m.group(1)
 159         msg = m.group(2)
 160
 161     return head, msg
 162
 163 def get_loglevel():
 164     return (k for k,v in LOG_LEVELS.items() if v==LOG_LEVEL).next()
 165
 166 def set_loglevel(level):
 167     global LOG_LEVEL
 168     if level not in LOG_LEVELS:
 169         # no effect
 170         return
 171
 172     LOG_LEVEL = LOG_LEVELS[level]
 173
 174 def set_interactive(mode=True):
 175     global INTERACTIVE
 176     if mode:
 177         INTERACTIVE = True
 178     else:
 179         INTERACTIVE = False
 180
 181 def raw(msg=''):
 182     _general_print('', NO_COLOR, msg)
 183
 184 def info(msg):
 185     head, msg = _split_msg('Info', msg)
 186     _general_print(head, INFO_COLOR, msg)
 187
 188 def verbose(msg):
 189     head, msg = _split_msg('Verbose', msg)
 190     _general_print(head, INFO_COLOR, msg, level = 'verbose')
 191
 192 def warning(msg):
 193     head, msg = _split_msg('Warning', msg)
 194     _color_perror(head, WARN_COLOR, msg)
 195
 196 def debug(msg):
 197     head, msg = _split_msg('Debug', msg)
 198     _color_perror(head, ERR_COLOR, msg, level = 'debug')
 199
 200 def error(msg):
 201     head, msg = _split_msg('Error', msg)
 202     _color_perror(head, ERR_COLOR, msg)
 203     sys.exit(1)
 204
 205 def ask(msg, default=True):
 206     _general_print('\rQ', ASK_COLOR, '')
 207     try:
 208         if default:
 209             msg += '(Y/n) '
 210         else:
 211             msg += '(y/N) '
 212         if INTERACTIVE:
 213             while True:
 214                 repl = raw_input(msg)
 215                 if repl.lower() == 'y':
 216                     return True
 217                 elif repl.lower() == 'n':
 218                     return False
 219                 elif not repl.strip():
 220                     # <Enter>
 221                     return default
 222
 223                 # else loop
 224         else:
 225             if default:
 226                 msg += ' Y'
 227             else:
 228                 msg += ' N'
 229             _general_print('', NO_COLOR, msg)
 230
 231             return default
 232     except KeyboardInterrupt:
 233         sys.stdout.write('\n')
 234         sys.exit(2)
 235
 236 def pause(msg=None):
 237     if INTERACTIVE:
 238         _general_print('\rQ', ASK_COLOR, '')
 239         if msg is None:
 240             msg = 'press <ENTER> to continue ...'
 241         raw_input(msg)
 242
 243 def set_logfile(fpath):
 244     global LOG_FILE_FP
 245
 246     def _savelogf():
 247         if LOG_FILE_FP:
 248             if not os.path.exists(os.path.dirname(LOG_FILE_FP)):
 249                 os.makedirs(os.path.dirname(LOG_FILE_FP))
 250             fp = open(LOG_FILE_FP, 'w')
 251             fp.write(LOG_CONTENT)
 252             fp.close()
 253
 254     if LOG_FILE_FP is not None:
 255         warning('duplicate log file configuration')
 256
 257     LOG_FILE_FP = os.path.abspath(os.path.expanduser(fpath))
 258
 259     import atexit
 260     atexit.register(_savelogf)
 261
 262 def enable_logstderr(fpath):
 263     global CATCHERR_BUFFILE_FD
 264     global CATCHERR_BUFFILE_PATH
 265     global CATCHERR_SAVED_2
 266
 267     if os.path.exists(fpath):
 268         os.remove(fpath)
 269     CATCHERR_BUFFILE_PATH = fpath
 270     CATCHERR_BUFFILE_FD = os.open(CATCHERR_BUFFILE_PATH, os.O_RDWR|os.O_CREAT)
 271     CATCHERR_SAVED_2 = os.dup(2)
 272     os.dup2(CATCHERR_BUFFILE_FD, 2)
 273
 274 def disable_logstderr():
 275     global CATCHERR_BUFFILE_FD
 276     global CATCHERR_BUFFILE_PATH
 277     global CATCHERR_SAVED_2
 278
 279     raw(msg = None) # flush message buffer and print it.
 280     os.dup2(CATCHERR_SAVED_2, 2)
 281     os.close(CATCHERR_SAVED_2)
 282     os.close(CATCHERR_BUFFILE_FD)
 283     os.unlink(CATCHERR_BUFFILE_PATH)
 284     CATCHERR_BUFFILE_FD = -1
 285     CATCHERR_BUFFILE_PATH = None
 286     CATCHERR_SAVED_2 = -1