1"""
2
3Helper functions for writing to terminals and files.
4
5"""
6
7
8import sys, os
9import py
10py3k = sys.version_info[0] >= 3
11from py.builtin import text, bytes
12
13win32_and_ctypes = False
14colorama = None
15if sys.platform == "win32":
16    try:
17        import colorama
18    except ImportError:
19        try:
20            import ctypes
21            win32_and_ctypes = True
22        except ImportError:
23            pass
24
25
26def _getdimensions():
27    import termios,fcntl,struct
28    call = fcntl.ioctl(1,termios.TIOCGWINSZ,"\000"*8)
29    height,width = struct.unpack( "hhhh", call ) [:2]
30    return height, width
31
32
33def get_terminal_width():
34    width = 0
35    try:
36        _, width = _getdimensions()
37    except py.builtin._sysex:
38        raise
39    except:
40        # pass to fallback below
41        pass
42
43    if width == 0:
44        # FALLBACK:
45        # * some exception happened
46        # * or this is emacs terminal which reports (0,0)
47        width = int(os.environ.get('COLUMNS', 80))
48
49    # XXX the windows getdimensions may be bogus, let's sanify a bit
50    if width < 40:
51        width = 80
52    return width
53
54terminal_width = get_terminal_width()
55
56# XXX unify with _escaped func below
57def ansi_print(text, esc, file=None, newline=True, flush=False):
58    if file is None:
59        file = sys.stderr
60    text = text.rstrip()
61    if esc and not isinstance(esc, tuple):
62        esc = (esc,)
63    if esc and sys.platform != "win32" and file.isatty():
64        text = (''.join(['\x1b[%sm' % cod for cod in esc])  +
65                text +
66                '\x1b[0m')     # ANSI color code "reset"
67    if newline:
68        text += '\n'
69
70    if esc and win32_and_ctypes and file.isatty():
71        if 1 in esc:
72            bold = True
73            esc = tuple([x for x in esc if x != 1])
74        else:
75            bold = False
76        esctable = {()   : FOREGROUND_WHITE,                 # normal
77                    (31,): FOREGROUND_RED,                   # red
78                    (32,): FOREGROUND_GREEN,                 # green
79                    (33,): FOREGROUND_GREEN|FOREGROUND_RED,  # yellow
80                    (34,): FOREGROUND_BLUE,                  # blue
81                    (35,): FOREGROUND_BLUE|FOREGROUND_RED,   # purple
82                    (36,): FOREGROUND_BLUE|FOREGROUND_GREEN, # cyan
83                    (37,): FOREGROUND_WHITE,                 # white
84                    (39,): FOREGROUND_WHITE,                 # reset
85                    }
86        attr = esctable.get(esc, FOREGROUND_WHITE)
87        if bold:
88            attr |= FOREGROUND_INTENSITY
89        STD_OUTPUT_HANDLE = -11
90        STD_ERROR_HANDLE = -12
91        if file is sys.stderr:
92            handle = GetStdHandle(STD_ERROR_HANDLE)
93        else:
94            handle = GetStdHandle(STD_OUTPUT_HANDLE)
95        oldcolors = GetConsoleInfo(handle).wAttributes
96        attr |= (oldcolors & 0x0f0)
97        SetConsoleTextAttribute(handle, attr)
98        while len(text) > 32768:
99            file.write(text[:32768])
100            text = text[32768:]
101        if text:
102            file.write(text)
103        SetConsoleTextAttribute(handle, oldcolors)
104    else:
105        file.write(text)
106
107    if flush:
108        file.flush()
109
110def should_do_markup(file):
111    if os.environ.get('PY_COLORS') == '1':
112        return True
113    if os.environ.get('PY_COLORS') == '0':
114        return False
115    return hasattr(file, 'isatty') and file.isatty() \
116           and os.environ.get('TERM') != 'dumb' \
117           and not (sys.platform.startswith('java') and os._name == 'nt')
118
119class TerminalWriter(object):
120    _esctable = dict(black=30, red=31, green=32, yellow=33,
121                     blue=34, purple=35, cyan=36, white=37,
122                     Black=40, Red=41, Green=42, Yellow=43,
123                     Blue=44, Purple=45, Cyan=46, White=47,
124                     bold=1, light=2, blink=5, invert=7)
125
126    # XXX deprecate stringio argument
127    def __init__(self, file=None, stringio=False, encoding=None):
128        if file is None:
129            if stringio:
130                self.stringio = file = py.io.TextIO()
131            else:
132                from sys import stdout as file
133        elif py.builtin.callable(file) and not (
134             hasattr(file, "write") and hasattr(file, "flush")):
135            file = WriteFile(file, encoding=encoding)
136        if hasattr(file, "isatty") and file.isatty() and colorama:
137            file = colorama.AnsiToWin32(file).stream
138        self.encoding = encoding or getattr(file, 'encoding', "utf-8")
139        self._file = file
140        self.hasmarkup = should_do_markup(file)
141        self._lastlen = 0
142        self._chars_on_current_line = 0
143
144    @property
145    def fullwidth(self):
146        if hasattr(self, '_terminal_width'):
147            return self._terminal_width
148        return get_terminal_width()
149
150    @fullwidth.setter
151    def fullwidth(self, value):
152        self._terminal_width = value
153
154    @property
155    def chars_on_current_line(self):
156        """Return the number of characters written so far in the current line.
157
158        Please note that this count does not produce correct results after a reline() call,
159        see #164.
160
161        .. versionadded:: 1.5.0
162
163        :rtype: int
164        """
165        return self._chars_on_current_line
166
167    def _escaped(self, text, esc):
168        if esc and self.hasmarkup:
169            text = (''.join(['\x1b[%sm' % cod for cod in esc])  +
170                text +'\x1b[0m')
171        return text
172
173    def markup(self, text, **kw):
174        esc = []
175        for name in kw:
176            if name not in self._esctable:
177                raise ValueError("unknown markup: %r" %(name,))
178            if kw[name]:
179                esc.append(self._esctable[name])
180        return self._escaped(text, tuple(esc))
181
182    def sep(self, sepchar, title=None, fullwidth=None, **kw):
183        if fullwidth is None:
184            fullwidth = self.fullwidth
185        # the goal is to have the line be as long as possible
186        # under the condition that len(line) <= fullwidth
187        if sys.platform == "win32":
188            # if we print in the last column on windows we are on a
189            # new line but there is no way to verify/neutralize this
190            # (we may not know the exact line width)
191            # so let's be defensive to avoid empty lines in the output
192            fullwidth -= 1
193        if title is not None:
194            # we want 2 + 2*len(fill) + len(title) <= fullwidth
195            # i.e.    2 + 2*len(sepchar)*N + len(title) <= fullwidth
196            #         2*len(sepchar)*N <= fullwidth - len(title) - 2
197            #         N <= (fullwidth - len(title) - 2) // (2*len(sepchar))
198            N = (fullwidth - len(title) - 2) // (2*len(sepchar))
199            fill = sepchar * N
200            line = "%s %s %s" % (fill, title, fill)
201        else:
202            # we want len(sepchar)*N <= fullwidth
203            # i.e.    N <= fullwidth // len(sepchar)
204            line = sepchar * (fullwidth // len(sepchar))
205        # in some situations there is room for an extra sepchar at the right,
206        # in particular if we consider that with a sepchar like "_ " the
207        # trailing space is not important at the end of the line
208        if len(line) + len(sepchar.rstrip()) <= fullwidth:
209            line += sepchar.rstrip()
210
211        self.line(line, **kw)
212
213    def write(self, msg, **kw):
214        if msg:
215            if not isinstance(msg, (bytes, text)):
216                msg = text(msg)
217
218            self._update_chars_on_current_line(msg)
219
220            if self.hasmarkup and kw:
221                markupmsg = self.markup(msg, **kw)
222            else:
223                markupmsg = msg
224            write_out(self._file, markupmsg)
225
226    def _update_chars_on_current_line(self, text):
227        fields = text.rsplit('\n', 1)
228        if '\n' in text:
229            self._chars_on_current_line = len(fields[-1])
230        else:
231            self._chars_on_current_line += len(fields[-1])
232
233    def line(self, s='', **kw):
234        self.write(s, **kw)
235        self._checkfill(s)
236        self.write('\n')
237
238    def reline(self, line, **kw):
239        if not self.hasmarkup:
240            raise ValueError("cannot use rewrite-line without terminal")
241        self.write(line, **kw)
242        self._checkfill(line)
243        self.write('\r')
244        self._lastlen = len(line)
245
246    def _checkfill(self, line):
247        diff2last = self._lastlen - len(line)
248        if diff2last > 0:
249            self.write(" " * diff2last)
250
251class Win32ConsoleWriter(TerminalWriter):
252    def write(self, msg, **kw):
253        if msg:
254            if not isinstance(msg, (bytes, text)):
255                msg = text(msg)
256
257            self._update_chars_on_current_line(msg)
258
259            oldcolors = None
260            if self.hasmarkup and kw:
261                handle = GetStdHandle(STD_OUTPUT_HANDLE)
262                oldcolors = GetConsoleInfo(handle).wAttributes
263                default_bg = oldcolors & 0x00F0
264                attr = default_bg
265                if kw.pop('bold', False):
266                    attr |= FOREGROUND_INTENSITY
267
268                if kw.pop('red', False):
269                    attr |= FOREGROUND_RED
270                elif kw.pop('blue', False):
271                    attr |= FOREGROUND_BLUE
272                elif kw.pop('green', False):
273                    attr |= FOREGROUND_GREEN
274                elif kw.pop('yellow', False):
275                    attr |= FOREGROUND_GREEN|FOREGROUND_RED
276                else:
277                    attr |= oldcolors & 0x0007
278
279                SetConsoleTextAttribute(handle, attr)
280            write_out(self._file, msg)
281            if oldcolors:
282                SetConsoleTextAttribute(handle, oldcolors)
283
284class WriteFile(object):
285    def __init__(self, writemethod, encoding=None):
286        self.encoding = encoding
287        self._writemethod = writemethod
288
289    def write(self, data):
290        if self.encoding:
291            data = data.encode(self.encoding, "replace")
292        self._writemethod(data)
293
294    def flush(self):
295        return
296
297
298if win32_and_ctypes:
299    TerminalWriter = Win32ConsoleWriter
300    import ctypes
301    from ctypes import wintypes
302
303    # ctypes access to the Windows console
304    STD_OUTPUT_HANDLE = -11
305    STD_ERROR_HANDLE  = -12
306    FOREGROUND_BLACK     = 0x0000 # black text
307    FOREGROUND_BLUE      = 0x0001 # text color contains blue.
308    FOREGROUND_GREEN     = 0x0002 # text color contains green.
309    FOREGROUND_RED       = 0x0004 # text color contains red.
310    FOREGROUND_WHITE     = 0x0007
311    FOREGROUND_INTENSITY = 0x0008 # text color is intensified.
312    BACKGROUND_BLACK     = 0x0000 # background color black
313    BACKGROUND_BLUE      = 0x0010 # background color contains blue.
314    BACKGROUND_GREEN     = 0x0020 # background color contains green.
315    BACKGROUND_RED       = 0x0040 # background color contains red.
316    BACKGROUND_WHITE     = 0x0070
317    BACKGROUND_INTENSITY = 0x0080 # background color is intensified.
318
319    SHORT = ctypes.c_short
320    class COORD(ctypes.Structure):
321        _fields_ = [('X', SHORT),
322                    ('Y', SHORT)]
323    class SMALL_RECT(ctypes.Structure):
324        _fields_ = [('Left', SHORT),
325                    ('Top', SHORT),
326                    ('Right', SHORT),
327                    ('Bottom', SHORT)]
328    class CONSOLE_SCREEN_BUFFER_INFO(ctypes.Structure):
329        _fields_ = [('dwSize', COORD),
330                    ('dwCursorPosition', COORD),
331                    ('wAttributes', wintypes.WORD),
332                    ('srWindow', SMALL_RECT),
333                    ('dwMaximumWindowSize', COORD)]
334
335    _GetStdHandle = ctypes.windll.kernel32.GetStdHandle
336    _GetStdHandle.argtypes = [wintypes.DWORD]
337    _GetStdHandle.restype = wintypes.HANDLE
338    def GetStdHandle(kind):
339        return _GetStdHandle(kind)
340
341    SetConsoleTextAttribute = ctypes.windll.kernel32.SetConsoleTextAttribute
342    SetConsoleTextAttribute.argtypes = [wintypes.HANDLE, wintypes.WORD]
343    SetConsoleTextAttribute.restype = wintypes.BOOL
344
345    _GetConsoleScreenBufferInfo = \
346        ctypes.windll.kernel32.GetConsoleScreenBufferInfo
347    _GetConsoleScreenBufferInfo.argtypes = [wintypes.HANDLE,
348                                ctypes.POINTER(CONSOLE_SCREEN_BUFFER_INFO)]
349    _GetConsoleScreenBufferInfo.restype = wintypes.BOOL
350    def GetConsoleInfo(handle):
351        info = CONSOLE_SCREEN_BUFFER_INFO()
352        _GetConsoleScreenBufferInfo(handle, ctypes.byref(info))
353        return info
354
355    def _getdimensions():
356        handle = GetStdHandle(STD_OUTPUT_HANDLE)
357        info = GetConsoleInfo(handle)
358        # Substract one from the width, otherwise the cursor wraps
359        # and the ending \n causes an empty line to display.
360        return info.dwSize.Y, info.dwSize.X - 1
361
362def write_out(fil, msg):
363    # XXX sometimes "msg" is of type bytes, sometimes text which
364    # complicates the situation.  Should we try to enforce unicode?
365    try:
366        # on py27 and above writing out to sys.stdout with an encoding
367        # should usually work for unicode messages (if the encoding is
368        # capable of it)
369        fil.write(msg)
370    except UnicodeEncodeError:
371        # on py26 it might not work because stdout expects bytes
372        if fil.encoding:
373            try:
374                fil.write(msg.encode(fil.encoding))
375            except UnicodeEncodeError:
376                # it might still fail if the encoding is not capable
377                pass
378            else:
379                fil.flush()
380                return
381        # fallback: escape all unicode characters
382        msg = msg.encode("unicode-escape").decode("ascii")
383        fil.write(msg)
384    fil.flush()
385