path: root/configshell/console.py

'''
This file is part of ConfigShell Community Edition.
Copyright (c) 2011 by RisingTide Systems LLC

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU Affero General Public License as
published by the Free Software Foundation, version 3 (AGPLv3).

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU Affero General Public License for more details.

You should have received a copy of the GNU Affero General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/>.
'''

import re
import sys
import tty
import prefs
import struct
import textwrap
from fcntl import ioctl
from termios import TIOCGWINSZ, TCSADRAIN, tcsetattr, tcgetattr

# avoid requiring epydoc at runtime
try:
    import epydoc.markup.epytext
except ImportError:
    pass

class Console(object):
    '''
    Implements various utility methods providing a console UI support toolkit,
    most notably an epytext-to-console text renderer using ANSI escape
    sequences. It uses the Borg pattern to share state between instances.
    '''
    _max_width = 132
    _escape = '\033['
    _ansi_format = _escape + '%dm%s'
    _ansi_reset = _escape + '0m'
    _re_ansi_seq = re.compile('(\033\[..?m)')

    _ansi_styles = {'bold':      1,
                    'underline': 4,
                    'blink':     5,
                    'reverse':   7,
                    'concealed': 8}

    colors = ['black', 'red', 'green', 'yellow',
              'blue', 'magenta', 'cyan', 'white']

    _ansi_fgcolors = dict(zip(colors, range(30, 38)))
    _ansi_bgcolors = dict(zip(colors, range(40, 48)))

    __borg_state = {}

    def __init__(self, stdin=sys.stdin, stdout=sys.stdout):
        '''
        Initializes a Console instance.
        @param stdin: The console standard input.
        @type stdin: file object
        @param stdout: The console standard output.
        @type stdout: file object
        '''
        self.__dict__ = self.__borg_state
        self._stdout = stdout
        self._stdin = stdin
        self.prefs = prefs.Prefs()

    # Public methods

    def escape(self, sequence, reply_terminator=None):
        '''
        Sends an escape sequence to the console, and reads the reply terminated
        by reply_terminator. If reply_terminator is not specified, the reply
        will not be read.
        @type sequence: str
        @param reply_terminator: The expected end-of-reply marker.
        @type reply_terminator: str
        '''
        attributes = tcgetattr(self._stdin)
        tty.setraw(self._stdin)
        try:
            self.raw_write(self._escape + sequence)
            if reply_terminator is not None:
                reply = ''
                while reply[-len(reply_terminator):] != reply_terminator:
                    reply += self._stdin.read(1)
        finally:
            tcsetattr(self._stdin, TCSADRAIN, attributes)
        if reply_terminator is not None:
            reply = reply[:-len(reply_terminator)]
            reply = reply.replace(self._escape, '').split(';')
            return reply

    def get_width(self):
        '''
        Returns the console width, or maximum width if we are not a terminal
        device.
        '''
        try:
            winsize = struct.pack("HHHH", 0, 0, 0, 0)
            winsize = ioctl(self._stdout.fileno(), TIOCGWINSZ, winsize)
            width = struct.unpack("HHHH", winsize)[1]
        except IOError:
            width = self._max_width
        else:
            if width > self._max_width:
                width = self._max_width

        return width

    def get_cursor_xy(self):
        '''
        Get the current text cursor x, y coordinates.
        '''
        coords = [int(coord) for coord in self.escape("6n", "R")]
        coords.reverse()
        return coords

    def set_cursor_xy(self, xpos, ypos):
        '''
        Set the cursor x, y coordinates.
        @param xpos: The x coordinate of the cursor.
        @type xpos: int
        @param ypos: The y coordinate of the cursor.
        @type ypos: int
        '''
        self.escape("%d;%dH" % (ypos, xpos))

    def raw_write(self, text):
        '''
        Raw console printing function.
        @param text: The text to print.
        @type text: str
        '''
        self._stdout.write(text)
        self._stdout.flush()

    def display(self, text, no_lf=False):
        '''
        Display a text with a default style.
        @param text: Text to display
        @type text: str
        @param no_lf: Do not display a line feed.
        @type no_lf: bool
        '''
        text = self.render_text(text)
        self.raw_write(text)
        if not no_lf:
            self.raw_write('\n')

    def epy_write(self, text):
        '''
        Renders and print and epytext-formatted text on the console.
        '''
        text = self.dedent(text)
        try:
            dom_tree = epydoc.markup.epytext.parse(text, None)
        except NameError:
            # epydoc not installed, strip markup
            dom_tree = text
            dom_tree = dom_tree.replace("B{", "")
            dom_tree = dom_tree.replace("I{", "")
            dom_tree = dom_tree.replace("C{", "")
            dom_tree = dom_tree.replace("}", "")
            dom_tree += "\n"
        except:
            self.display(text)
            raise
        text = self.render_domtree(dom_tree)
        # We need to remove the last line feed, but there might be
        # escape characters after it...
        clean_text = ''
        for index in range(1, len(text)):
            if text[-index] == '\n':
                clean_text = text[:-index]
                if index != 1:
                    clean_text += text[-index+1:]
                break
        else:
            clean_text = text
        self.raw_write(clean_text)

    def indent(self, text, margin=2):
        '''
        Indents text by margin space.
        @param text: The text to be indented.
        @type text: str
        '''
        output = ''
        for line in text.split('\n'):
            output += margin * ' ' + line + '\n'
        return output

    def dedent(self, text):
        '''
        A convenience function to easily write multiline text blocks that
        will be later assembled in to a unique epytext string.
        It removes heading newline chars and common indentation.
        '''
        for i in range(len(text)):
            if text[i] != '\n':
                break
        text = text[i:]
        text = textwrap.dedent(text)
        text = '\n' * i + text

        return text

    def render_text(self, text, fgcolor=None, bgcolor=None, styles=None,
                    open_end=False, todefault=False):
        '''
        Renders some text with ANSI console colors and attributes.
        @param fgcolor: ANSI color to use for text:
            black, red, green, yellow, blue, magenta. cyan. white
        @type fgcolor: str
        @param bgcolor: ANSI color to use for background:
            black, red, green, yellow, blue, magenta. cyan. white
        @type bgcolor: str
        @param styles: List of ANSI styles to use:
            bold, underline, blink, reverse, concealed
        @type styles: list of str
        @param open_end: Do not reset text style at the end ot the output.
        @type open_end: bool
        @param todefault: Instead of resetting style at the end of the
        output, reset to default color. Only if not open_end.
        @type todefault: bool
        '''
        if self.prefs['color_mode'] and self._stdout.isatty():
            if fgcolor is None:
                if self.prefs['color_default']:
                    fgcolor = self.prefs['color_default']
            if fgcolor is not None:
                text = self._ansi_format % (self._ansi_fgcolors[fgcolor], text)
            if bgcolor is not None:
                text = self._ansi_format % (self._ansi_bgcolors[bgcolor], text)
            if styles is not None:
                for style in styles:
                    text = self._ansi_format % (self._ansi_styles[style], text)
            if not open_end:
                text += self._ansi_reset
                if todefault and fgcolor is not None:
                    if self.prefs['color_default']:
                        text += self._ansi_format \
                                % (self._ansi_fgcolors[
                                    self.prefs['color_default']], '')
        return text

    def wordwrap(self, text, indent=0, startindex=0, splitchars=''):
        '''
        Word-wrap the given string.  I.e., add newlines to the string such
        that any lines that are longer than terminal width or max_width
        are broken into shorter lines (at the first whitespace sequence that
        occurs before the limit. If the given string contains newlines, they
        will I{not} be removed.  Any lines that begin with whitespace will not
        be wordwrapped.

        This version takes into account ANSI escape characters:
            - stop escape sequence styling at the end of a split line
            - start it again on the next line if needed after the indent
            - do not account for the length of the escape sequences when
              wrapping

        @param indent: If specified, then indent each line by this number
            of spaces.
        @type indent: C{int}
        @param startindex: If specified, then assume that the first line
            is already preceeded by C{startindex} characters.
        @type startindex: C{int}
        @param splitchars: A list of non-whitespace characters which can
            be used to split a line.  (E.g., use '/\\' to allow path names
            to be split over multiple lines.)
        @rtype: C{str}
        '''
        right = self.get_width()
        if splitchars:
            chunks = re.split(r'( +|\n|[^ \n%s]*[%s])' %
                              (re.escape(splitchars), re.escape(splitchars)),
                              text.expandtabs())
        else:
            chunks = re.split(r'( +|\n)', text.expandtabs())
        result = [' '*(indent-startindex)]
        charindex = max(indent, startindex)
        current_style = ''
        for chunknum, chunk in enumerate(chunks):
            chunk_groups = re.split(self._re_ansi_seq, chunk)
            chunk_text = ''
            next_style = current_style

            for group in chunk_groups:
                if re.match(self._re_ansi_seq, group) is None:
                    chunk_text += group
                else:
                    next_style += group

            chunk_len = len(chunk_text)
            if (charindex + chunk_len > right and charindex > 0) \
               or chunk == '\n':
                result[-1] = result[-1].rstrip()
                result.append(self.render_text(
                    '\n' + ' '*indent + current_style, open_end=True))
                charindex = indent
                if chunk[:1] not in ('\n', ' '):
                    result.append(chunk)
                    charindex += chunk_len
            else:
                result.append(chunk)
                charindex += chunk_len

            current_style = next_style.split(self._ansi_reset)[-1]

        return ''.join(result).rstrip()+'\n'

    def render_domtree(self, tree, indent=0, seclevel=0):
        '''
        Convert a DOM document encoding epytext to an 8-bits ascii string with
        ANSI formating for simpler styles.

        @param tree: A DOM document encoding of an epytext string.
        @type tree: C{Element}
        @param indent: The indentation for the string representation of
            C{tree}.  Each line of the returned string will begin with
            C{indent} space characters.
        @type indent: C{int}
        @param seclevel: The section level that C{tree} appears at.  This
            is used to generate section headings.
        @type seclevel: C{int}
        @return: The formated string.
        @rtype: C{string}
        '''
        if isinstance(tree, basestring):
            return tree

        if tree.tag == 'section':
            seclevel += 1

        # Figure out the child indent level.
        if tree.tag == 'epytext':
            cindent = indent
        elif tree.tag == 'li' and tree.attribs.get('bullet'):
            cindent = indent + 1 + len(tree.attribs.get('bullet'))
        else:
            cindent = indent + 2

        variables = [self.render_domtree(c, cindent, seclevel)
                     for c in tree.children]
        childstr = ''.join(variables)

        if tree.tag == 'para':
            text = self.render_text(childstr)
            text = self.wordwrap(text, indent)+'\n'
        elif tree.tag == 'li':
            # We should be able to use getAttribute here; but there's no
            # convenient way to test if an element has an attribute..
            bullet = tree.attribs.get('bullet') or '-'
            text = indent*' ' + bullet + ' ' + childstr.lstrip()
        elif tree.tag == 'heading':
            text = ((indent-2)*' ' + self.render_text(
                childstr, styles=['bold'], todefault=True) \
                + '\n')
        elif tree.tag == 'doctestblock':
            lines = [(indent+2)*' '+line for line in childstr.split('\n')]
            text = '\n'.join(lines) + '\n\n'
        elif tree.tag == 'literalblock':
            lines = [(indent+1)*' '+ self.render_text(
                line, todefault=True)
                     for line in childstr.split('\n')]
            text = '\n'.join(lines) + '\n\n'
        elif tree.tag == 'fieldlist':
            text = childstr
        elif tree.tag == 'field':
            numargs = 0
            while tree.children[numargs+1].tag == 'arg':
                numargs += 1
            args = variables[1:1+numargs]
            body = variables[1+numargs:]
            text = (indent)*' '+'@'+variables[0]
            if args:
                text += '(' + ', '.join(args) + ')'
            text = text + ':\n' + ''.join(body)
        elif tree.tag == 'uri':
            if len(variables) != 2:
                raise ValueError('Bad URI ')
            elif variables[0] == variables[1]:
                text = self.render_text(
                    '%s' % variables[1],
                    'blue', styles=['underline'], todefault=True)
            else:
                text = '%r<%s>' % (variables[0], variables[1])
        elif tree.tag == 'link':
            if len(variables) != 2:
                raise ValueError('Bad Link')
            text = '%s' % variables[0]
        elif tree.tag in ('olist', 'ulist'):
            text = childstr.replace('\n\n', '\n')+'\n'
        elif tree.tag == 'bold':
            text = self.render_text(
                childstr, styles=['bold'], todefault=True)
        elif tree.tag == 'italic':
            text = self.render_text(
                childstr, styles=['underline'], todefault=True)
        elif tree.tag == 'symbol':
            text = '%s' \
                    % epydoc.markup.epytext.SYMBOL_TO_PLAINTEXT.get(
                        childstr, childstr)
        elif tree.tag == 'graph':
            text = '<<%s graph: %s>>' \
                    % (variables[0], ', '.join(variables[1:]))
        else:
            # Assume that anything else can be passed through.
            text = self.render_text(childstr)

        return text