1# Testoob, Python Testing Out Of (The) Box 2# Copyright (C) 2005-2006 The Testoob Team 3# 4# Licensed under the Apache License, Version 2.0 (the "License"); 5# you may not use this file except in compliance with the License. 6# You may obtain a copy of the License at 7# 8# http://www.apache.org/licenses/LICENSE-2.0 9# 10# Unless required by applicable law or agreed to in writing, software 11# distributed under the License is distributed on an "AS IS" BASIS, 12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13# See the License for the specific language governing permissions and 14# limitations under the License. 15 16"Color text stream reporting" 17 18import os, sys 19 20ANSI_CODES = { 21 "reset" : "\x1b[0m", 22 "bold" : "\x1b[01m", 23 "teal" : "\x1b[36;06m", 24 "turquoise" : "\x1b[36;01m", 25 "fuscia" : "\x1b[35;01m", 26 "purple" : "\x1b[35;06m", 27 "blue" : "\x1b[34;01m", 28 "darkblue" : "\x1b[34;06m", 29 "green" : "\x1b[32;01m", 30 "darkgreen" : "\x1b[32;06m", 31 "yellow" : "\x1b[33;01m", 32 "brown" : "\x1b[33;06m", 33 "red" : "\x1b[31;01m", 34} 35 36from textstream import StreamWriter 37class TerminalColorWriter(StreamWriter): 38 def __init__(self, stream, color): 39 StreamWriter.__init__(self, stream) 40 self.code = ANSI_CODES[color] 41 self.reset = ANSI_CODES["reset"] 42 def write(self, s): 43 StreamWriter.write(self, self.code) 44 StreamWriter.write(self, s) 45 StreamWriter.write(self, self.reset) 46 def get_bgcolor(self): 47 return "unknown" 48 49 50class WindowsColorBaseWriter(StreamWriter): 51 """ 52 All Windows writers set the color without writing special control 53 characters, so this class is convenient. 54 """ 55 FOREGROUND_BLUE = 0x0001 # text color contains blue. 56 FOREGROUND_GREEN = 0x0002 # text color contains green. 57 FOREGROUND_RED = 0x0004 # text color contains red. 58 FOREGROUND_INTENSITY = 0x0008 # text color is intensified. 59 BACKGROUND_BLUE = 0x0010 # background color contains blue. 60 BACKGROUND_GREEN = 0x0020 # background color contains green. 61 BACKGROUND_RED = 0x0040 # background color contains red. 62 BACKGROUND_INTENSITY = 0x0080 # background color is intensified. 63 64 def __init__(self, stream, color): 65 StreamWriter.__init__(self, stream) 66 self.reset = self._get_color() 67 self.background = self.reset & 0xf0 68 CODES = { 69 "red" : self.FOREGROUND_RED | self.FOREGROUND_INTENSITY | self.background, 70 "green" : self.FOREGROUND_GREEN | self.FOREGROUND_INTENSITY | self.background, 71 "yellow" : self.FOREGROUND_GREEN | self.FOREGROUND_RED | self.FOREGROUND_INTENSITY | self.background, 72 "blue" : self.FOREGROUND_BLUE | self.FOREGROUND_INTENSITY | self.background 73 } 74 self.code = CODES[color] 75 76 def write(self, s): 77 self._set_color(self.code) 78 StreamWriter.write(self, s) 79 self._set_color(self.reset) 80 81 def get_bgcolor(self): 82 WHITE = self.BACKGROUND_RED | self.BACKGROUND_GREEN | self.BACKGROUND_BLUE | self.BACKGROUND_INTENSITY 83 YELLOW = self.BACKGROUND_RED | self.BACKGROUND_GREEN | self.BACKGROUND_INTENSITY 84 if self.background in [WHITE, YELLOW]: 85 return "light" 86 else: 87 return "dark" 88 89 90class Win32ColorWriterWithExecutable(WindowsColorBaseWriter): 91 setcolor_path = os.path.join(sys.prefix, "testoob", "setcolor.exe") 92 setcolor_available = os.path.isfile(setcolor_path) 93 if not setcolor_available: 94 setcolor_path = os.path.join('other','setcolor.exe') 95 setcolor_available = os.path.isfile(setcolor_path) 96 97 def _set_color(self, code): 98 # TODO: fail in advance if setcolor.exe isn't available? 99 if self.setcolor_available: 100 try: 101 import subprocess 102 except ImportError: 103 from testoob.compatibility import subprocess 104 subprocess.Popen('"%s" set %d' % (self.setcolor_path, code)).wait() 105 106 def _get_color(self): 107 if self.setcolor_available: 108 try: 109 import subprocess 110 except ImportError: 111 from testoob.compatibility import subprocess 112 get_pipe = subprocess.Popen('"%s" get' % (self.setcolor_path), 113 stdout=subprocess.PIPE) 114 color_code, _ = get_pipe.communicate() 115 return int(color_code) 116 else: 117 return 0x0f 118 119 120class Win32ConsoleColorWriter(WindowsColorBaseWriter): 121 def _out_handle(self): 122 import win32console 123 return win32console.GetStdHandle(win32console.STD_OUTPUT_HANDLE) 124 out_handle = property(_out_handle) 125 126 def _set_color(self, code): 127 self.out_handle.SetConsoleTextAttribute( code ) 128 129 def _get_color(self): 130 return self.out_handle.GetConsoleScreenBufferInfo()['Attributes'] 131 132 133class WindowsCtypesColorWriter(WindowsColorBaseWriter): 134 # Constants from the Windows API 135 STD_OUTPUT_HANDLE = -11 136 137 def _out_handle(self): 138 import ctypes 139 return ctypes.windll.kernel32.GetStdHandle(self.STD_OUTPUT_HANDLE) 140 out_handle = property(_out_handle) 141 142 def _console_screen_buffer_info(self): 143 # Based on IPython's winconsole.py, written by Alexander Belchenko 144 import ctypes, struct 145 csbi = ctypes.create_string_buffer(22) 146 res = ctypes.windll.kernel32.GetConsoleScreenBufferInfo(self.out_handle, csbi) 147 assert res 148 149 (bufx, bufy, curx, cury, wattr, 150 left, top, right, bottom, maxx, maxy) = struct.unpack("hhhhHhhhhhh", csbi.raw) 151 152 return { 153 "bufx" : bufx, 154 "bufy" : bufy, 155 "curx" : curx, 156 "cury" : cury, 157 "wattr" : wattr, 158 "left" : left, 159 "top" : top, 160 "right" : right, 161 "bottom" : bottom, 162 "maxx" : maxx, 163 "maxy" : maxy, 164 } 165 console_screen_buffer_info = property(_console_screen_buffer_info) 166 167 def _set_color(self, code): 168 import ctypes 169 ctypes.windll.kernel32.SetConsoleTextAttribute(self.out_handle, code) 170 171 def _get_color(self): 172 return self.console_screen_buffer_info["wattr"] 173 174 175def color_writers_creator(writer_class): 176 class ColorWriters: 177 def _get_warning_color(self, bgcolor): 178 import options 179 if options.bgcolor != "auto": 180 bgcolor = options.bgcolor 181 bg_mapping = {"dark": "yellow", "light": "blue", "unknown": "yellow" } 182 warning_color = bg_mapping[bgcolor] 183 return warning_color 184 185 def __init__(self, stream): 186 self.normal = StreamWriter(stream) 187 self.success = writer_class(stream, "green") 188 self.failure = writer_class(stream, "red") 189 bgcolor = self.success.get_bgcolor() 190 self.warning = writer_class(stream, self._get_warning_color(bgcolor)) 191 return ColorWriters 192 193from textstream import TextStreamReporter 194def create_colored_reporter(writer_class): 195 class ColoredReporter(TextStreamReporter): 196 def __init__(self, *args, **kwargs): 197 kwargs["create_writers"] = color_writers_creator(writer_class) 198 TextStreamReporter.__init__(self, *args, **kwargs) 199 return ColoredReporter 200 201def choose_color_writer(): 202 if "TESTOOB_COLOR_WRITER" in os.environ: 203 #print "DEBUG: using", os.environ["TESTOOB_COLOR_WRITER"] 204 return eval(os.environ["TESTOOB_COLOR_WRITER"]) 205 206 if sys.platform != "win32": 207 return TerminalColorWriter 208 209 try: 210 import win32console 211 return Win32ConsoleColorWriter 212 except ImportError: 213 pass 214 215 try: 216 import ctypes 217 return WindowsCtypesColorWriter 218 except ImportError: 219 pass 220 221 return Win32ColorWriterWithExecutable 222 223ColoredTextReporter = create_colored_reporter( choose_color_writer() ) 224