highlighting.py

Go to the documentation of this file.

#  Copyright 2008-2015 Nokia Networks
#  Copyright 2016-     Robot Framework Foundation
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
 
# Windows highlighting code adapted from color_console.py. It is copyright
# Andre Burgaud, licensed under the MIT License, and available here:
# http://www.burgaud.com/bring-colors-to-the-windows-console-with-python/
 
from contextlib import contextmanager
import errno
import os
import sys
try:
    from ctypes import windll, Structure, c_short, c_ushort, byref
except ImportError:  # Not on Windows
    windll = None
 
from robot.errors import DataError
from robot.utils import console_encode, isatty, WINDOWS
 
 
class HighlightingStream:
 
    def __init__(self, stream, colors='AUTO'):
        self.streamstream = stream
        self._highlighter_highlighter = self._get_highlighter_get_highlighter(stream, colors)
 
    def _get_highlighter(self, stream, colors):
        options = {'AUTO': Highlighter if isatty(stream) else NoHighlighting,
                   'ON': Highlighter,
                   'OFF': NoHighlighting,
                   'ANSI': AnsiHighlighter}
        try:
            highlighter = options[colors.upper()]
        except KeyError:
            raise DataError("Invalid console color value '%s'. Available "
                            "'AUTO', 'ON', 'OFF' and 'ANSI'." % colors)
        return highlighter(stream)
 
    def write(self, text, flush=True):
        self._write_write(console_encode(text, stream=self.streamstream))
        if flush:
            self.flushflush()
 
    def _write(self, text, retry=5):
        # Workaround for Windows 10 console bug:
        # https://github.com/robotframework/robotframework/issues/2709
        try:
            with self._suppress_broken_pipe_error_suppress_broken_pipe_error_suppress_broken_pipe_error:
                self.streamstream.write(text)
        except IOError as err:
            if not (WINDOWS and err.errno == 0 and retry > 0):
                raise
            self._write_write(text, retry-1)
 
    @property
    @contextmanager
    _suppress_broken_pipe_error = property
    
    def _suppress_broken_pipe_error(self):
        try:
            yield
        except IOError as err:
            if err.errno not in (errno.EPIPE, errno.EINVAL, errno.EBADF):
                raise
 
    def flush(self):
        with self._suppress_broken_pipe_error_suppress_broken_pipe_error_suppress_broken_pipe_error:
            self.streamstream.flush()
 
    def highlight(self, text, status=None, flush=True):
        if self._must_flush_before_and_after_highlighting_must_flush_before_and_after_highlighting():
            self.flushflush()
            flush = True
        with self._highlighting_highlighting(status or text):
            self.writewrite(text, flush)
 
    def _must_flush_before_and_after_highlighting(self):
        # Must flush on Windows before and after highlighting to make sure set
        # console colors only affect the actual highlighted text. Problems
        # only encountered with Python 3, but better to be safe than sorry.
        return WINDOWS and not isinstance(self._highlighter_highlighter, NoHighlighting)
 
    def error(self, message, level):
        self.writewrite('[ ', flush=False)
        self.highlighthighlight(level, flush=False)
        self.writewrite(' ] %s\n' % message)
 
    @contextmanager
    def _highlighting(self, status):
        highlighter = self._highlighter_highlighter
        start = {'PASS': highlighter.green,
                 'FAIL': highlighter.red,
                 'ERROR': highlighter.red,
                 'WARN': highlighter.yellow,
                 'SKIP': highlighter.yellow}[status]
        start()
        try:
            yield
        finally:
            highlighter.reset()
 
 
def Highlighter(stream):
    if os.sep == '/':
        return AnsiHighlighter(stream)
    return DosHighlighter(stream) if windll else NoHighlighting(stream)
 
 
class AnsiHighlighter:
    
    _ANSI_GREEN = '\033[32m'
    
    _ANSI_RED = '\033[31m'
    
    _ANSI_YELLOW = '\033[33m'
    
    _ANSI_RESET = '\033[0m'
 
    def __init__(self, stream):
        self._stream_stream = stream
 
    def green(self):
        self._set_color_set_color(self._ANSI_GREEN_ANSI_GREEN)
 
    def red(self):
        self._set_color_set_color(self._ANSI_RED_ANSI_RED)
 
    def yellow(self):
        self._set_color_set_color(self._ANSI_YELLOW_ANSI_YELLOW)
 
    def reset(self):
        self._set_color_set_color(self._ANSI_RESET_ANSI_RESET)
 
    def _set_color(self, color):
        self._stream_stream.write(color)
 
 
class NoHighlighting(AnsiHighlighter):
 
    def _set_color(self, color):
        pass
 
 
class DosHighlighter:
    
    _FOREGROUND_GREEN = 0x2
    
    _FOREGROUND_RED = 0x4
    
    _FOREGROUND_YELLOW = 0x6
    
    _FOREGROUND_GREY = 0x7
    
    _FOREGROUND_INTENSITY = 0x8
    
    _BACKGROUND_MASK = 0xF0
    
    _STDOUT_HANDLE = -11
    
    _STDERR_HANDLE = -12
 
    def __init__(self, stream):
        self._handle_handle = self._get_std_handle_get_std_handle(stream)
        self._orig_colors_orig_colors = self._get_colors_get_colors()
        self._background_background = self._orig_colors_orig_colors & self._BACKGROUND_MASK_BACKGROUND_MASK
 
    def green(self):
        self._set_foreground_colors_set_foreground_colors(self._FOREGROUND_GREEN_FOREGROUND_GREEN)
 
    def red(self):
        self._set_foreground_colors_set_foreground_colors(self._FOREGROUND_RED_FOREGROUND_RED)
 
    def yellow(self):
        self._set_foreground_colors_set_foreground_colors(self._FOREGROUND_YELLOW_FOREGROUND_YELLOW)
 
    def reset(self):
        self._set_colors_set_colors(self._orig_colors_orig_colors)
 
    def _get_std_handle(self, stream):
        handle = self._STDOUT_HANDLE_STDOUT_HANDLE \
            if stream is sys.__stdout__ else self._STDERR_HANDLE_STDERR_HANDLE
        return windll.kernel32.GetStdHandle(handle)
 
    def _get_colors(self):
        csbi = _CONSOLE_SCREEN_BUFFER_INFO()
        ok = windll.kernel32.GetConsoleScreenBufferInfo(self._handle_handle, byref(csbi))
        if not ok:  # Call failed, return default console colors (gray on black)
            return self._FOREGROUND_GREY_FOREGROUND_GREY
        return csbi.wAttributes
 
    def _set_foreground_colors(self, colors):
        self._set_colors_set_colors(colors | self._FOREGROUND_INTENSITY_FOREGROUND_INTENSITY | self._background_background)
 
    def _set_colors(self, colors):
        windll.kernel32.SetConsoleTextAttribute(self._handle_handle, colors)
 
 
if windll:
 
    class _COORD(Structure):
        
        _fields_ = [("X", c_short),
                    ("Y", c_short)]
 
    class _SMALL_RECT(Structure):
        
        _fields_ = [("Left", c_short),
                    ("Top", c_short),
                    ("Right", c_short),
                    ("Bottom", c_short)]
 
    class _CONSOLE_SCREEN_BUFFER_INFO(Structure):
        
        _fields_ = [("dwSize", _COORD),
                    ("dwCursorPosition", _COORD),
                    ("wAttributes", c_ushort),
                    ("srWindow", _SMALL_RECT),
                    ("dwMaximumWindowSize", _COORD)]