1 """

     2 

     3 Helper functions for writing to terminals and files.

     4 

     5 """

     6 

     7 

     8 import sys, os

     9 import py

    10 

    11 win32_and_ctypes = False

    12 if sys.platform == "win32":

    13     try:

    14         import ctypes

    15         win32_and_ctypes = True

    16     except ImportError:

    17         pass

    18 

    19 def _getdimensions():

    20     import termios,fcntl,struct

    21     call = fcntl.ioctl(1,termios.TIOCGWINSZ,"\000"*8)

    22     height,width = struct.unpack( "hhhh", call ) [:2]

    23     return height, width

    24 

    25 

    26 def get_terminal_width():

    27     try:

    28         height, width = _getdimensions()

    29     except py.builtin._sysex:

    30         raise

    31     except:

    32         # FALLBACK

    33         width = int(os.environ.get('COLUMNS', 80))

    34     else:

    35         # XXX the windows getdimensions may be bogus, let's sanify a bit

    36         if width < 40:

    37             width = 80

    38     return width

    39 

    40 terminal_width = get_terminal_width()

    41 

    42 # XXX unify with _escaped func below

    43 def ansi_print(text, esc, file=None, newline=True, flush=False):

    44     if file is None:

    45         file = sys.stderr

    46     text = text.rstrip()

    47     if esc and not isinstance(esc, tuple):

    48         esc = (esc,)

    49     if esc and sys.platform != "win32" and file.isatty():

    50         text = (''.join(['\x1b[%sm' % cod for cod in esc])  +

    51                 text +

    52                 '\x1b[0m')     # ANSI color code "reset"

    53     if newline:

    54         text += '\n'

    55 

    56     if esc and win32_and_ctypes and file.isatty():

    57         if 1 in esc:

    58             bold = True

    59             esc = tuple([x for x in esc if x != 1])

    60         else:

    61             bold = False

    62         esctable = {()   : FOREGROUND_WHITE,                 # normal

    63                     (31,): FOREGROUND_RED,                   # red

    64                     (32,): FOREGROUND_GREEN,                 # green

    65                     (33,): FOREGROUND_GREEN|FOREGROUND_RED,  # yellow

    66                     (34,): FOREGROUND_BLUE,                  # blue

    67                     (35,): FOREGROUND_BLUE|FOREGROUND_RED,   # purple

    68                     (36,): FOREGROUND_BLUE|FOREGROUND_GREEN, # cyan

    69                     (37,): FOREGROUND_WHITE,                 # white

    70                     (39,): FOREGROUND_WHITE,                 # reset

    71                     }

    72         attr = esctable.get(esc, FOREGROUND_WHITE)

    73         if bold:

    74             attr |= FOREGROUND_INTENSITY

    75         STD_OUTPUT_HANDLE = -11

    76         STD_ERROR_HANDLE = -12

    77         if file is sys.stderr:

    78             handle = GetStdHandle(STD_ERROR_HANDLE)

    79         else:

    80             handle = GetStdHandle(STD_OUTPUT_HANDLE)

    81         oldcolors = GetConsoleInfo(handle).wAttributes

    82         attr |= (oldcolors & 0x0f0)

    83         SetConsoleTextAttribute(handle, attr)

    84         file.write(text)

    85         SetConsoleTextAttribute(handle, oldcolors)

    86     else:

    87         file.write(text)

    88 

    89     if flush:

    90         file.flush()

    91 

    92 def should_do_markup(file):

    93     return hasattr(file, 'isatty') and file.isatty() \

    94            and os.environ.get('TERM') != 'dumb' \

    95            and not (sys.platform.startswith('java') and os._name == 'nt')

    96 

    97 class TerminalWriter(object):

    98     _esctable = dict(black=30, red=31, green=32, yellow=33,

    99                      blue=34, purple=35, cyan=36, white=37,

   100                      Black=40, Red=41, Green=42, Yellow=43,

   101                      Blue=44, Purple=45, Cyan=46, White=47,

   102                      bold=1, light=2, blink=5, invert=7)

   103 

   104     # XXX deprecate stringio argument

   105     def __init__(self, file=None, stringio=False, encoding=None):

   106         if file is None:

   107             if stringio:

   108                 self.stringio = file = py.io.TextIO()

   109             else:

   110                 file = py.std.sys.stdout

   111                 if hasattr(file, 'encoding'):

   112                     encoding = file.encoding

   113         elif hasattr(file, '__call__'):

   114             file = WriteFile(file, encoding=encoding)

   115         self.encoding = encoding

   116         self._file = file

   117         self.fullwidth = get_terminal_width()

   118         self.hasmarkup = should_do_markup(file)

   119 

   120     def _escaped(self, text, esc):

   121         if esc and self.hasmarkup:

   122             text = (''.join(['\x1b[%sm' % cod for cod in esc])  +

   123                 text +'\x1b[0m')

   124         return text

   125 

   126     def markup(self, text, **kw):

   127         esc = []

   128         for name in kw:

   129             if name not in self._esctable:

   130                 raise ValueError("unknown markup: %r" %(name,))

   131             if kw[name]:

   132                 esc.append(self._esctable[name])

   133         return self._escaped(text, tuple(esc))

   134 

   135     def sep(self, sepchar, title=None, fullwidth=None, **kw):

   136         if fullwidth is None:

   137             fullwidth = self.fullwidth

   138         # the goal is to have the line be as long as possible

   139         # under the condition that len(line) <= fullwidth

   140         if title is not None:

   141             # we want 2 + 2*len(fill) + len(title) <= fullwidth

   142             # i.e.    2 + 2*len(sepchar)*N + len(title) <= fullwidth

   143             #         2*len(sepchar)*N <= fullwidth - len(title) - 2

   144             #         N <= (fullwidth - len(title) - 2) // (2*len(sepchar))

   145             N = (fullwidth - len(title) - 2) // (2*len(sepchar))

   146             fill = sepchar * N

   147             line = "%s %s %s" % (fill, title, fill)

   148         else:

   149             # we want len(sepchar)*N <= fullwidth

   150             # i.e.    N <= fullwidth // len(sepchar)

   151             line = sepchar * (fullwidth // len(sepchar))

   152         # in some situations there is room for an extra sepchar at the right,

   153         # in particular if we consider that with a sepchar like "_ " the

   154         # trailing space is not important at the end of the line

   155         if len(line) + len(sepchar.rstrip()) <= fullwidth:

   156             line += sepchar.rstrip()

   157 

   158         self.line(line, **kw)

   159 

   160     def write(self, s, **kw):

   161         if s:

   162             if not isinstance(self._file, WriteFile):

   163                 s = self._getbytestring(s)

   164                 if self.hasmarkup and kw:

   165                     s = self.markup(s, **kw)

   166             self._file.write(s)

   167             self._file.flush()

   168 

   169     def _getbytestring(self, s):

   170         # XXX review this and the whole logic

   171         if self.encoding and sys.version_info[0] < 3 and isinstance(s, unicode):

   172             return s.encode(self.encoding)

   173         elif not isinstance(s, str):

   174             try:

   175                 return str(s)

   176             except UnicodeEncodeError:

   177                 return "<print-error '%s' object>" % type(s).__name__

   178         return s

   179 

   180     def line(self, s='', **kw):

   181         self.write(s, **kw)

   182         self.write('\n')

   183 

   184 class Win32ConsoleWriter(TerminalWriter):

   185     def write(self, s, **kw):

   186         if s:

   187             oldcolors = None

   188             if self.hasmarkup and kw:

   189                 handle = GetStdHandle(STD_OUTPUT_HANDLE)

   190                 oldcolors = GetConsoleInfo(handle).wAttributes

   191                 default_bg = oldcolors & 0x00F0

   192                 attr = default_bg

   193                 if kw.pop('bold', False):

   194                     attr |= FOREGROUND_INTENSITY

   195 

   196                 if kw.pop('red', False):

   197                     attr |= FOREGROUND_RED

   198                 elif kw.pop('blue', False):

   199                     attr |= FOREGROUND_BLUE

   200                 elif kw.pop('green', False):

   201                     attr |= FOREGROUND_GREEN

   202                 else:

   203                     attr |= FOREGROUND_BLACK # (oldcolors & 0x0007)

   204 

   205                 SetConsoleTextAttribute(handle, attr)

   206             if not isinstance(self._file, WriteFile):

   207                 s = self._getbytestring(s)

   208             self._file.write(s)

   209             self._file.flush()

   210             if oldcolors:

   211                 SetConsoleTextAttribute(handle, oldcolors)

   212 

   213     def line(self, s="", **kw):

   214         self.write(s+"\n", **kw)

   215 

   216 class WriteFile(object):

   217     def __init__(self, writemethod, encoding=None):

   218         self.encoding = encoding

   219         self._writemethod = writemethod

   220 

   221     def write(self, data):

   222         if self.encoding:

   223             data = data.encode(self.encoding)

   224         self._writemethod(data)

   225 

   226     def flush(self):

   227         return

   228 

   229 

   230 if win32_and_ctypes:

   231     TerminalWriter = Win32ConsoleWriter

   232     import ctypes

   233     from ctypes import wintypes

   234 

   235     # ctypes access to the Windows console

   236     STD_OUTPUT_HANDLE = -11

   237     STD_ERROR_HANDLE  = -12

   238     FOREGROUND_BLACK     = 0x0000 # black text

   239     FOREGROUND_BLUE      = 0x0001 # text color contains blue.

   240     FOREGROUND_GREEN     = 0x0002 # text color contains green.

   241     FOREGROUND_RED       = 0x0004 # text color contains red.

   242     FOREGROUND_WHITE     = 0x0007

   243     FOREGROUND_INTENSITY = 0x0008 # text color is intensified.

   244     BACKGROUND_BLACK     = 0x0000 # background color black

   245     BACKGROUND_BLUE      = 0x0010 # background color contains blue.

   246     BACKGROUND_GREEN     = 0x0020 # background color contains green.

   247     BACKGROUND_RED       = 0x0040 # background color contains red.

   248     BACKGROUND_WHITE     = 0x0070

   249     BACKGROUND_INTENSITY = 0x0080 # background color is intensified.

   250 

   251     SHORT = ctypes.c_short

   252     class COORD(ctypes.Structure):

   253         _fields_ = [('X', SHORT),

   254                     ('Y', SHORT)]

   255     class SMALL_RECT(ctypes.Structure):

   256         _fields_ = [('Left', SHORT),

   257                     ('Top', SHORT),

   258                     ('Right', SHORT),

   259                     ('Bottom', SHORT)]

   260     class CONSOLE_SCREEN_BUFFER_INFO(ctypes.Structure):

   261         _fields_ = [('dwSize', COORD),

   262                     ('dwCursorPosition', COORD),

   263                     ('wAttributes', wintypes.WORD),

   264                     ('srWindow', SMALL_RECT),

   265                     ('dwMaximumWindowSize', COORD)]

   266 

   267     def GetStdHandle(kind):

   268         return ctypes.windll.kernel32.GetStdHandle(kind)

   269 

   270     SetConsoleTextAttribute = \

   271         ctypes.windll.kernel32.SetConsoleTextAttribute

   272 

   273     def GetConsoleInfo(handle):

   274         info = CONSOLE_SCREEN_BUFFER_INFO()

   275         ctypes.windll.kernel32.GetConsoleScreenBufferInfo(\

   276             handle, ctypes.byref(info))

   277         return info

   278 

   279     def _getdimensions():

   280         handle = GetStdHandle(STD_OUTPUT_HANDLE)

   281         info = GetConsoleInfo(handle)

   282         # Substract one from the width, otherwise the cursor wraps

   283         # and the ending \n causes an empty line to display.

   284         return info.dwSize.Y, info.dwSize.X - 1

   285