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