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