scripteditor.py

"""Qt script editor for Python with syntax highlighting.

Author: Peter Hunt
Updated: 5/9/2019
"""

try:
    import builtins
except ImportError:
    import __builtin__ as builtins
import keyword
from Qt import QtCore, QtGui, QtWidgets


class ScriptEditor(QtWidgets.QTextEdit):
    """A QTextEdit with key overrides to act like a script editor."""
    def __init__(self, *args, **kwargs):
        super(ScriptEditor, self).__init__(*args, **kwargs)
        self.setTabSize(4)

    def tabSize(self):
        return self._tabSize

    def setTabSize(self, tabSize):
        self._tabSize = tabSize
        self._tabIndent = ' ' * tabSize

    def tabIndent(self):
        return self._tabIndent

    def setTabIndent(self, indent):
        self.setTabSize(indent.count(' ') + indent.count('\t') * 4)

    def keyPressEvent(self, event):
        key = event.key()
        if key in (QtCore.Qt.Key_Tab, QtCore.Qt.Key_Backtab):
            return self.__tabPressEvent(event)
        elif key == QtCore.Qt.Key_Return:
            return self.__returnPressEvent(event)
        elif key == QtCore.Qt.Key_Backspace:
            return self.__backspacePressEvent(event)
        return super(ScriptEditor, self).keyPressEvent(event)

    def __tabPressEvent(self, event):
        # Catch cases where backtab isn't automatically set
        key = event.key()
        if key == QtCore.Qt.Key_Tab and event.modifiers() & QtCore.Qt.ShiftModifier:
            key = QtCore.Qt.Key_Backtab

        text = self.toPlainText()

        # Get the selection
        textCursor = self.textCursor()
        selectionStart = textCursor.selectionStart()
        selectionEnd = textCursor.selectionEnd()

        # Add a tab to the current cursor
        if selectionStart == selectionEnd and key == QtCore.Qt.Key_Tab:
            self.setPlainText('{}{}{}'.format(text[:selectionStart], self.tabIndent(), text[selectionEnd:]))
            textCursor.setPosition(selectionStart + self.tabSize())
            self.setTextCursor(textCursor)

        # Move selected text
        else:
            # Expand selection to start and end of lines
            textCursor.setPosition(selectionStart)
            textCursor.movePosition(QtGui.QTextCursor.StartOfLine, QtGui.QTextCursor.KeepAnchor)
            textCursor.setPosition(textCursor.position())
            textCursor.setPosition(selectionEnd, QtGui.QTextCursor.KeepAnchor)
            textCursor.movePosition(QtGui.QTextCursor.EndOfLine, QtGui.QTextCursor.KeepAnchor)

            # Get the selection from the start to end of lines
            lineSelectionStart = textCursor.selectionStart()
            lineSelectionEnd = textCursor.selectionEnd()
            
            # Split the selected text into individual lines
            middleText = text[lineSelectionStart:lineSelectionEnd].split('\n')
            numLines = len(filter(bool, middleText))
            
            # Edit each line and work out the new indexes of the selection
            if key == QtCore.Qt.Key_Tab:
                middleText = [self.tabIndent() + line if line else line for line in middleText]
                if selectionStart != lineSelectionStart:
                    selectionStart += self.tabSize()
                selectionEnd += self.tabSize() * numLines
                
            elif key == QtCore.Qt.Key_Backtab:
                if not all(line[:self.tabSize()] == self.tabIndent() if line.strip() else True for line in middleText):
                    return
                middleText = [line[self.tabSize():] for line in middleText]
                if selectionStart != lineSelectionStart:
                    selectionStart -= self.tabSize()
                selectionEnd -= self.tabSize() * numLines
            
            self.setPlainText(text[:lineSelectionStart] + '\n'.join(middleText) + text[lineSelectionEnd:])

            # Update cursor
            textCursor.setPosition(selectionStart)
            textCursor.setPosition(selectionEnd, QtGui.QTextCursor.KeepAnchor)
            self.setTextCursor(textCursor)

    def __returnPressEvent(self, event):
        # Get the selection
        textCursor = self.textCursor()
        selectionStart = textCursor.selectionStart()
        if not selectionStart:
            return super(ScriptEditor, self).keyPressEvent(event)
        text = self.toPlainText()

        # Find if the line ends with a colon
        index = selectionStart - 1
        while text[index] == ' ':
            index -= 1
        indents = bool(text[index] == ':')

        # Find how many indents are at the start of the string
        textCursor.movePosition(QtGui.QTextCursor.StartOfLine, QtGui.QTextCursor.MoveAnchor)
        numSpaces = 0
        index = textCursor.position()
        try:
            while text[numSpaces+index] == ' ':
                numSpaces += 1
        except IndexError:
            pass
        indents += numSpaces // self.tabSize()

        # Calculate new text
        insertSpaces = self.tabIndent() * indents
        newText = text[:selectionStart] + '\n' + insertSpaces + text[selectionStart:]
        self.setPlainText(newText)

        # Calculate new cursor position
        textCursor.setPosition(selectionStart + self.tabSize()*indents + 1, QtGui.QTextCursor.MoveAnchor)
        self.setTextCursor(textCursor)

    def __backspacePressEvent(self, event):
        textCursor = self.textCursor()
        selectionStart = textCursor.selectionStart()
        selectionEnd = textCursor.selectionEnd()

        # If there is a selection, then just remove that normally
        if selectionStart != selectionEnd:
            return super(ScriptEditor, self).keyPressEvent(event)

        # If the last <tabSize> characters are not all spaces, treat normally
        text = self.toPlainText()
        if text[selectionStart-self.tabSize():selectionStart] != self.tabIndent():
            return super(ScriptEditor, self).keyPressEvent(event)

        newText = text[:selectionStart-self.tabSize()] + text[selectionStart:]
        self.setPlainText(newText)
        textCursor.setPosition(selectionStart-self.tabSize(), QtGui.QTextCursor.MoveAnchor)
        self.setTextCursor(textCursor)


def formatColour(colorName, italic=False, bold=False):
    """Return a QtGui.QTextCharFormat with the given attributes.
    
    Source: https://wiki.python.org/moin/PyQt/Python%20syntax%20highlighting
    """
    colour = QtGui.QColor()
    colour.setNamedColor(colorName)

    textFormat = QtGui.QTextCharFormat()
    textFormat.setForeground(colour)
    if bold:
        textFormat.setFontWeight(QtGui.QFont.Bold)
    if italic:
        textFormat.setFontItalic(True)

    return textFormat


class PythonHighlighter(QtGui.QSyntaxHighlighter):
    """Base Python highlighter class.
    This must be subclassed and will not work by itself.
    
    Source: https://wiki.python.org/moin/PyQt/Python%20syntax%20highlighting

    Required Attribute Overrides:
        Colours:
            Dictionary of {'key': formatColour('colour')}
            Possible keys:
                keyword: Things like "and", "as", "break", "class", "if"
                builtin: Things like "eval, "file", "round", "set", "str"
                comment: Single line comments
                string: Strings with single or double quotes
                docstring: Multi line comments
                defclass: Special case for def and class
    """
    Keywords = keyword.kwlist

    Builtins = dir(builtins)

    Operators = [
        '=',
        # Comparison
        '==', '!=', '<', '<=', '>', '>=',
        # Arithmetic
        '\+', '-', '\*', '/', '//', '\%', '\*\*',
        # In-place
        '\+=', '-=', '\*=', '/=', '\%=',
        # Bitwise
        '\^', '\|', '\&', '\~', '>>', '<<',
    ]

    Colours = {}
    
    def __init__(self, parent, **kwargs):
        super(PythonHighlighter, self).__init__(parent)

        # Multi-line strings (expression, flag, style)
        if 'docstring' in self.Colours:
            self.tri_single = (QtCore.QRegExp("'''"), 1, self.Colours['docstring'])
            self.tri_double = (QtCore.QRegExp('"""'), 2, self.Colours['docstring'])
        
        rules = []

        if 'keyword' in self.Colours:
            rules += [(r'\b{}\b'.format(w), 0, self.Colours['keyword']) for w in self.Keywords]
        
        if 'operator' in self.Colours:
            rules += [(r'\b{}\b'.format(w), 0, self.Colours['operator']) for w in self.Operators]
        
        if 'builtin' in self.Colours:
            rules += [(r'\b{}\b'.format(w), 0, self.Colours['builtin']) for w in self.Builtins]
        
        if 'string' in self.Colours:
            # Double-quoted string, possibly containing escape sequences
            rules.append((r'"[^"\\]*(\\.[^"\\]*)*"', 0, self.Colours['string']))
            # Single-quoted string, possibly containing escape sequences
            rules.append((r"'[^'\\]*(\\.[^'\\]*)*'", 0, self.Colours['string']))
        
        # TODO: Avoid this if it is in a string
        if 'defclass' in self.Colours:
            # 'def' followed by an identifier
            rules.append((r'\bdef\b\s*(\w+)', 1, self.Colours['defclass']))
            # 'def' followed by an identifier
            rules.append((r'\bclass\b\s*(\w+)', 1, self.Colours['defclass']))

        if 'comment' in self.Colours:
            rules.append((r'#[^\n]*', 0, self.Colours['comment']))

        if 'numbers' in self.Colours:
            rules.append((r'\b[+-]?[0-9]+[lL]?\b', 0, self.Colours['numbers'])),
            rules.append((r'\b[+-]?0[xX][0-9A-Fa-f]+[lL]?\b', 0, self.Colours['numbers'])),
            rules.append((r'\b[+-]?[0-9]+(?:\.[0-9]+)?(?:[eE][+-]?[0-9]+)?\b', 0, self.Colours['numbers'])),
        
        if 'maya' in self.Colours:
            try:
                import maya.cmds as cmds
            except ImportError:
                pass
            else:
                commands = cmds.help('[a-z]*', list=True, lng='Python')
                rules += [(r'\b{}\b'.format(w), 0, self.Colours['maya']) for w in commands]
        
        # Build a QRegExp for each pattern
        self.rules = [(QtCore.QRegExp(pat), index, fmt)
            for (pat, index, fmt) in rules]

    def highlightBlock(self, text):
        """Apply syntax highlighting to the given block of text."""
        # Do other syntax formatting
        for expression, nth, format in self.rules:
            index = expression.indexIn(text, 0)

            while index >= 0:
                # We actually want the index of the nth match
                index = expression.pos(nth)
                length = len(expression.cap(nth))
                self.setFormat(index, length, format)
                index = expression.indexIn(text, index + length)

        self.setCurrentBlockState(0)

        # Do multi-line strings
        if not self.match_multiline(text, *self.tri_single):
            self.match_multiline(text, *self.tri_double)

    def match_multiline(self, text, delimiter, in_state, style):
        # If inside triple-single quotes, start at 0
        if self.previousBlockState() == in_state:
            start = 0
            add = 0
        # Otherwise, look for the delimiter on this line
        else:
            start = delimiter.indexIn(text)
            # Move past this match
            add = delimiter.matchedLength()

        # As long as there's a delimiter match on this line...
        while start >= 0:
            # Look for the ending delimiter
            end = delimiter.indexIn(text, start + add)
            # Ending delimiter on this line?
            if end >= add:
                length = end - start + add + delimiter.matchedLength()
                self.setCurrentBlockState(0)
            # No; multi-line string
            else:
                self.setCurrentBlockState(in_state)
                length = len(text) - start + add
            # Apply formatting
            self.setFormat(start, length, style)
            # Look for the next match
            start = delimiter.indexIn(text, start + length)

        # Return True if still inside a multi-line string, False otherwise
        if self.currentBlockState() == in_state:
            return True
        else:
            return False


class PythonIDLEHighlighter(PythonHighlighter):
    """IDLE Python syntax highlighter."""
    Colours = {
        'keyword': formatColour('darkOrange'),
        'operator': formatColour('red'),
        'comment': formatColour('red'),
        'string': formatColour('green'),
        'docstring': formatColour('green'),
        'builtin': formatColour('darkMagenta'),
        'defclass': formatColour('blue'),
    }


class PythonMayaHighlighter(PythonHighlighter):
    """Maya Python syntax highlighter."""
    Colours = {
        'keyword': formatColour('lime'),
        'argument': formatColour('black'),
        'comment': formatColour('red'),
        'string': formatColour('yellow'),
        'docstring': formatColour('yellow', italic=True),
        'maya': formatColour('cyan')
    }


class PythonMASHHighlighter(PythonHighlighter):
    """Maya's MASH Python syntax highlighter.

    Source: C:\Program Files\Autodesk\Maya2018\plug-ins\MASH\scripts\MASH\syntax.py
    """
    Colours = {
        'keyword': formatColour('DeepSkyBlue'),
        'builtin': formatColour('DeepSkyBlue'),
        'operator': formatColour('DeepPink'),
        'brace': formatColour('darkGray'),
        'defclass': formatColour('MistyRose'),
        'string': formatColour('red'),
        'docstring': formatColour('yellow'),
        'comment': formatColour('Gray'),
        'self': formatColour('Plum'),
        'numbers': formatColour('GhostWhite'),
        'maya': formatColour('SpringGreen'),
    }


if __name__ == '__main__':
    import sys
    app = QtWidgets.QApplication(sys.argv)
    mainWin = ScriptEditor()
    PythonIDLEHighlighter(mainWin)
    mainWin.show()
    sys.exit(app.exec_())