Revision 23269c5b qa/qa_utils.py

b/qa/qa_utils.py
21 21
"""
22 22

  
23 23
import os
24
import sys
24 25
import subprocess
25 26

  
26 27
from ganeti import utils
......
29 30
import qa_error
30 31

  
31 32

  
33
_INFO_SEQ = None
34
_WARNING_SEQ = None
35
_ERROR_SEQ = None
36
_RESET_SEQ = None
37

  
38

  
39
def _SetupColours():
40
  """Initializes the colour constants.
41

  
42
  """
43
  global _INFO_SEQ, _WARNING_SEQ, _ERROR_SEQ, _RESET_SEQ
44

  
45
  try:
46
    import curses
47
  except ImportError:
48
    # Don't use colours if curses module can't be imported
49
    return
50

  
51
  curses.setupterm()
52

  
53
  _RESET_SEQ = curses.tigetstr("op")
54

  
55
  setaf = curses.tigetstr("setaf")
56
  _INFO_SEQ = curses.tparm(setaf, curses.COLOR_GREEN)
57
  _WARNING_SEQ = curses.tparm(setaf, curses.COLOR_YELLOW)
58
  _ERROR_SEQ = curses.tparm(setaf, curses.COLOR_RED)
59

  
60

  
61
_SetupColours()
62

  
63

  
32 64
def AssertEqual(first, second, msg=None):
33 65
  """Raises an error when values aren't equal.
34 66

  
......
113 145
  AssertEqual(p.wait(), 0)
114 146

  
115 147
  return p.stdout.read().strip()
148

  
149

  
150
def _PrintWithColor(text, seq):
151
  f = sys.stdout
152

  
153
  if not f.isatty():
154
    seq = None
155

  
156
  if seq:
157
    f.write(seq)
158

  
159
  f.write(text)
160
  f.write("\n")
161

  
162
  if seq:
163
    f.write(_RESET_SEQ)
164

  
165

  
166
def PrintWarning(text):
167
  return _PrintWithColor(text, _WARNING_SEQ)
168

  
169

  
170
def PrintError(f, text):
171
  return _PrintWithColor(text, _ERROR_SEQ)
172

  
173

  
174
def PrintInfo(f, text):
175
  return _PrintWithColor(text, _INFO_SEQ)

Also available in: Unified diff