eric6/UI/PythonDisViewer.py

# -*- coding: utf-8 -*-

# Copyright (c) 2020 Detlev Offenbach <detlev@die-offenbachs.de>
#

"""
Module implementing a widget to visualize the Python Disassembly for some
Python sources.
"""

import os
import dis

import enum


from PyQt5.QtCore import pyqtSlot, Qt, QTimer
from PyQt5.QtGui import QCursor, QBrush
from PyQt5.QtWidgets import (
    QTreeWidget, QApplication, QTreeWidgetItem, QAbstractItemView, QWidget,
    QVBoxLayout, QLabel, QMenu
)

import Preferences


class PythonDisViewerModes(enum.Enum):
    """
    Class implementing the disassembly viewer operation modes.
    """
    SourceDisassemblyMode = 0
    TracebackMode = 1


class PythonDisViewer(QWidget):
    """
    Class implementing a widget to visualize the Python Disassembly for some
    Python sources.
    """
    StartLineRole = Qt.UserRole
    EndLineRole = Qt.UserRole + 1
    
    def __init__(self, viewmanager,
                 mode=PythonDisViewerModes.SourceDisassemblyMode,
                 parent=None):
        """
        Constructor
        
        @param viewmanager reference to the viewmanager object
        @type ViewManager
        @param mode operation mode of the viewer
        @type int
        @param parent reference to the parent widget
        @type QWidget
        """
        super(PythonDisViewer, self).__init__(parent)
        
        self.__layout = QVBoxLayout(self)
        self.setLayout(self.__layout)
        self.__disWidget = QTreeWidget(self)
        self.__layout.addWidget(self.__disWidget)
        self.__layout.setContentsMargins(0, 0, 0, 0)
        
        self.__currentInfoLabel = QLabel(self.tr(
            "italic: current instruction"))
        self.__labeledInfoLabel = QLabel(self.tr(
            "bold: labeled instruction"))
        self.__layout.addWidget(self.__currentInfoLabel)
        self.__layout.addWidget(self.__labeledInfoLabel)
        
        self.setWindowTitle(self.tr("Disassembly"))
        
        self.__vm = viewmanager
        self.__vmConnected = False
        
        self.__mode = mode
        
        self.__editor = None
        self.__source = ""
        
        self.__disWidget.setHeaderLabels(
            [self.tr("Line"), self.tr("Offset"), self.tr("Operation"),
             self.tr("Parameters"), self.tr("Interpreted Parameters")])
        self.__disWidget.setSortingEnabled(False)
        self.__disWidget.setSelectionBehavior(QAbstractItemView.SelectRows)
        self.__disWidget.setSelectionMode(QAbstractItemView.SingleSelection)
        self.__disWidget.setAlternatingRowColors(True)
        
        self.__menu = QMenu(self.__disWidget)
        self.__menu.addAction(self.tr('Expand All'), self.__expandAll)
        self.__menu.addAction(self.tr('Collapse All'), self.__collapseAll)
        
        self.__disWidget.setContextMenuPolicy(Qt.CustomContextMenu)
        self.__disWidget.customContextMenuRequested.connect(
            self.__contextMenuRequested)
        
        self.__errorColor = QBrush(
            Preferences.getPython("DisViewerErrorColor"))
        self.__currentInstructionColor = QBrush(
            Preferences.getPython("DisViewerCurrentColor"))
        self.__jumpTargetColor = QBrush(
            Preferences.getPython("DisViewerLabeledColor"))
        
        self.__disWidget.itemClicked.connect(self.__disItemClicked)
        self.__disWidget.itemCollapsed.connect(self.__resizeColumns)
        self.__disWidget.itemExpanded.connect(self.__resizeColumns)
        
        if self.__mode == PythonDisViewerModes.SourceDisassemblyMode:
            self.__vm.disViewerStateChanged.connect(
                self.__disViewerStateChanged)
        
            self.hide()
        elif self.__mode == PythonDisViewerModes.TracebackMode:
            self.__styleLabels()
    
    def __contextMenuRequested(self, coord):
        """
        Private slot to show the context menu.
        
        @param coord position of the mouse pointer
        @type QPoint
        """
        coord = self.__disWidget.mapToGlobal(coord)
        self.__menu.popup(coord)
    
    def __editorChanged(self, editor):
        """
        Private slot to handle a change of the current editor.
        
        @param editor reference to the current editor
        @type Editor
        """
        if editor is not self.__editor:
            if self.__editor:
                self.__editor.clearAllHighlights()
            self.__editor = editor
            if self.__editor:
                self.__loadDIS()
    
    def __editorSaved(self, editor):
        """
        Private slot to reload the Disassembly after the connected editor was
        saved.
        
        @param editor reference to the editor that performed a save action
        @type Editor
        """
        if editor and editor is self.__editor:
            self.__loadDIS()
    
    def __editorLineChanged(self, editor, lineno):
        """
        Private slot to handle a mouse button double click in the editor.
        
        @param editor reference to the editor, that emitted the signal
        @type Editor
        @param lineno line number of the editor's cursor (zero based)
        @type int
        """
        if editor is self.__editor:
            if editor.isModified():
                # reload the source
                QTimer.singleShot(0, self.__loadDIS)
            
            # highlight the corresponding entry
            QTimer.singleShot(0, self.__selectItemForEditorLine)
    
    def __editorLanguageChanged(self, editor):
        """
        Private slot to handle a change of the editor language.
        
        @param editor reference to the editor which changed language
        @type Editor
        """
        if editor is self.__editor:
            QTimer.singleShot(0, self.__loadDIS)
    
    def __lastEditorClosed(self):
        """
        Private slot to handle the last editor closed signal of the view
        manager.
        """
        self.hide()
    
    def show(self):
        """
        Public slot to show the DIS viewer.
        """
        super(PythonDisViewer, self).show()
        
        if (
            self.__mode == PythonDisViewerModes.SourceDisassemblyMode and
            not self.__vmConnected
        ):
            self.__vm.editorChangedEd.connect(self.__editorChanged)
            self.__vm.editorSavedEd.connect(self.__editorSaved)
            self.__vm.editorLineChangedEd.connect(self.__editorLineChanged)
            self.__vm.editorLanguageChanged.connect(
                self.__editorLanguageChanged)
            self.__vmConnected = True
        
        self.__styleLabels()
    
    def hide(self):
        """
        Public slot to hide the DIS viewer.
        """
        super(PythonDisViewer, self).hide()
        
        if self.__editor:
            self.__editor.clearAllHighlights()
        
        if (
            self.__mode == PythonDisViewerModes.SourceDisassemblyMode and
            self.__vmConnected
        ):
            self.__vm.editorChangedEd.disconnect(self.__editorChanged)
            self.__vm.editorSavedEd.disconnect(self.__editorSaved)
            self.__vm.editorLineChangedEd.disconnect(self.__editorLineChanged)
            self.__vm.editorLanguageChanged.disconnect(
                self.__editorLanguageChanged)
            self.__vmConnected = False
    
    def shutdown(self):
        """
        Public method to perform shutdown actions.
        """
        self.__editor = None
    
    def __disViewerStateChanged(self, on):
        """
        Private slot to toggle the display of the Disassembly viewer.
        
        @param on flag indicating to show the Disassembly
        @type bool
        """
        if self.__mode == PythonDisViewerModes.SourceDisassemblyMode:
            editor = self.__vm.activeWindow()
            if on:
                if editor is not self.__editor:
                    self.__editor = editor
                self.show()
                self.__loadDIS()
            else:
                self.hide()
                self.__editor = None
    
    def __expandAll(self):
        """
        Private slot to expand all items.
        """
        block = self.__disWidget.blockSignals(True)
        self.__disWidget.expandAll()
        self.__disWidget.blockSignals(block)
        self.__resizeColumns()
    
    def __collapseAll(self):
        """
        Private slot to collapse all items.
        """
        block = self.__disWidget.blockSignals(True)
        self.__disWidget.collapseAll()
        self.__disWidget.blockSignals(block)
        self.__resizeColumns()
    
    def __createErrorItem(self, error):
        """
        Private method to create a top level error item.
        
        @param error error message
        @type str
        @return generated item
        @rtype QTreeWidgetItem
        """
        itm = QTreeWidgetItem(self.__disWidget, [error])
        itm.setFirstColumnSpanned(True)
        itm.setForeground(0, self.__errorColor)
        return itm
    
    def __createTitleItem(self, title, line, parentItem):
        """
        Private method to create a title item.
        
        @param title titel string for the item
        @type str
        @param line start line of the titled disassembly
        @type int
        @param parentItem reference to the parent item
        @type QTreeWidget or QTreeWidgetItem
        @return generated item
        @rtype QTreeWidgetItem
        """
        itm = QTreeWidgetItem(parentItem, [title])
        itm.setFirstColumnSpanned(True)
        itm.setExpanded(True)
        
        itm.setData(0, self.StartLineRole, line)
        itm.setData(0, self.EndLineRole, line)
        
        return itm
    
    def __createInstructionItem(self, instr, parent, lasti=-1):
        """
        Private method to create an item for the given instruction.
        
        @param instr instruction the item should be based on
        @type dis.Instruction
        @param parent reference to the parent item
        @type QTreeWidgetItem
        @param lasti index of the instruction of a traceback
        @type int
        @return generated item
        @rtype QTreeWidgetItem
        """
        fields = []
        # Column: Source code line number (right aligned)
        if instr.starts_line:
            fields.append("{0:d}".format(instr.starts_line))
        else:
            fields.append("")
        # Column: Instruction offset from start of code sequence
        # (right aligned)
        fields.append("{0:d}".format(instr.offset))
        # Column: Opcode name
        fields.append(instr.opname)
        # Column: Opcode argument (right aligned)
        if instr.arg is not None:
            fields.append(repr(instr.arg))
            # Column: Opcode argument details
            if instr.argrepr:
                fields.append('(' + instr.argrepr + ')')
        
        itm = QTreeWidgetItem(parent, fields)
        for col in (0, 1, 3):
            itm.setTextAlignment(col, Qt.AlignRight)
        # set font to indicate current instruction and jump target
        font = itm.font(0)
        if instr.offset == lasti:
            font.setItalic(True)
        if instr.is_jump_target:
            font.setBold(True)
        for col in range(itm.columnCount()):
            itm.setFont(col, font)
        # set color to indicate current instruction or jump target
        if instr.offset == lasti:
            foreground = self.__currentInstructionColor
        elif instr.is_jump_target:
            foreground = self.__jumpTargetColor
        else:
            foreground = None
        if foreground:
            for col in range(itm.columnCount()):
                itm.setForeground(col, foreground)
        
        itm.setExpanded(True)
        
        if instr.starts_line:
            itm.setData(0, self.StartLineRole, instr.starts_line)
            itm.setData(0, self.EndLineRole, instr.starts_line)
        else:
            # get line from parent (= start line)
            lineno = parent.data(0, self.StartLineRole)
            itm.setData(0, self.StartLineRole, lineno)
            itm.setData(0, self.EndLineRole, lineno)
        return itm
    
    def __updateItemEndLine(self, itm):
        """
        Private method to update an items end line based on its children.
        
        @param itm reference to the item to be updated
        @type QTreeWidgetItem
        """
        if itm.childCount():
            endLine = max(
                itm.child(index).data(0, self.EndLineRole)
                for index in range(itm.childCount())
            )
        else:
            endLine = itm.data(0, self.StartLineRole)
        itm.setData(0, self.EndLineRole, endLine)
    
    def __loadDIS(self):
        """
        Private method to generate the Disassembly from the source of the
        current editor and visualize it.
        """
        if self.__mode != PythonDisViewerModes.SourceDisassemblyMode:
            # wrong mode, just return
            return
        
        if not self.__editor:
            self.__createErrorItem(self.tr(
                "No editor has been opened."
            ))
            return
        
        self.__disWidget.clear()
        self.__editor.clearAllHighlights()
        
        source = self.__editor.text()
        if not source.strip():
            # empty editor or white space only
            self.__createErrorItem(self.tr(
                "The current editor does not contain any source code."
            ))
            return
        
        if not self.__editor.isPyFile():
            self.__createErrorItem(self.tr(
                "The current editor does not contain Python source code."
            ))
            return
        
        filename = self.__editor.getFileName()
        if filename:
            filename = os.path.basename(filename)
        else:
            filename = "<dis>"
        
        QApplication.setOverrideCursor(QCursor(Qt.WaitCursor))
        try:
            codeObject = self.__tryCompile(source, filename)
        except Exception as exc:
            codeObject = None
            self.__createErrorItem(str(exc))
        
        if codeObject:
            self.setUpdatesEnabled(False)
            block = self.__disWidget.blockSignals(True)
            
            self.__disassembleObject(codeObject, self.__disWidget, filename)
            QTimer.singleShot(0, self.__resizeColumns)
            
            self.__disWidget.blockSignals(block)
            self.setUpdatesEnabled(True)
        
        QApplication.restoreOverrideCursor()
    
    @pyqtSlot(dict)
    def showDisassembly(self, disassembly):
        """
        Public slot to receive a code disassembly from the debug client.
        
        @param disassembly dictionary containing the disassembly information
        @type dict
        """
        if self.__mode == PythonDisViewerModes.TracebackMode:
            if (
                disassembly and
                "instructions" in disassembly and
                disassembly["instructions"]
            ):
                self.__disWidget.clear()
                
                self.setUpdatesEnabled(False)
                block = self.__disWidget.blockSignals(True)
                
                titleItem = self.__createTitleItem(
                    self.tr("Disassembly of last traceback"),
                    disassembly["firstlineno"],
                    self.__disWidget
                )
                
                lasti = disassembly["lasti"]
                lastStartItem = None
                for instrDict in disassembly["instructions"]:
                    instr = dis.Instruction(
                        instrDict["opname"],
                        0,                              # dummy value
                        instrDict["arg"],
                        "",                             # dummy value
                        instrDict["argrepr"],
                        instrDict["offset"],
                        instrDict["lineno"],
                        instrDict["isJumpTarget"],
                    )
                    if instrDict["lineno"] > 0:
                        if lastStartItem:
                            self.__updateItemEndLine(lastStartItem)
                        lastStartItem = self.__createInstructionItem(
                            instr, titleItem, lasti=lasti)
                    else:
                        self.__createInstructionItem(
                            instr, lastStartItem, lasti=lasti)
                if lastStartItem:
                    self.__updateItemEndLine(lastStartItem)
                
                QTimer.singleShot(0, self.__resizeColumns)
                
                self.__disWidget.blockSignals(block)
                self.setUpdatesEnabled(True)
                
                if lasti:
                    lastInstructions = self.__disWidget.findItems(
                        "{0:d}".format(lasti),
                        Qt.MatchFixedString | Qt.MatchRecursive, 1)
                    if lastInstructions:
                        self.__disWidget.scrollToItem(
                            lastInstructions[0],
                            QAbstractItemView.PositionAtCenter)
    
    def __resizeColumns(self):
        """
        Private method to resize the columns to suitable values.
        """
        for col in range(self.__disWidget.columnCount()):
            self.__disWidget.resizeColumnToContents(col)
    
    def resizeEvent(self, evt):
        """
        Protected method to handle resize events.
        
        @param evt resize event
        @type QResizeEvent
        """
        # just adjust the sizes of the columns
        self.__resizeColumns()
    
    def __clearSelection(self):
        """
        Private method to clear all selected items.
        """
        for itm in self.__disWidget.selectedItems():
            itm.setSelected(False)
    
    def __selectChildren(self, itm, lineno):
        """
        Private method to select children of the given item covering the given
        line number.
        
        @param itm reference to the item
        @type QTreeWidgetItem
        @param lineno line number to base the selection on
        @type int
        """
        for index in range(itm.childCount()):
            child = itm.child(index)
            if (
                child.data(0, self.StartLineRole) <= lineno <=
                child.data(0, self.EndLineRole)
            ):
                child.setSelected(True)
                self.__selectChildren(child, lineno)
            
            if child.data(0, self.StartLineRole) == lineno:
                self.__disWidget.scrollToItem(
                    child, QAbstractItemView.PositionAtCenter)
    
    def __selectItemForEditorLine(self):
        """
        Private slot to select the items corresponding with the cursor line
        of the current editor.
        """
        # step 1: clear all selected items
        self.__clearSelection()
        
        # step 2: retrieve the editor cursor line
        cline, cindex = self.__editor.getCursorPosition()
        # make the line numbers 1-based
        cline += 1
        
        for index in range(self.__disWidget.topLevelItemCount()):
            itm = self.__disWidget.topLevelItem(index)
            if (
                itm.data(0, self.StartLineRole) <= cline <=
                itm.data(0, self.EndLineRole)
            ):
                itm.setSelected(True)
                self.__selectChildren(itm, cline)
    
    @pyqtSlot(QTreeWidgetItem, int)
    def __disItemClicked(self, itm, column):
        """
        Private slot handling a user click on a Disassembly node item.
        
        @param itm reference to the clicked item
        @type QTreeWidgetItem
        @param column column number of the click
        @type int
        """
        # TODO: add code to deal with Traceback mode
        self.__editor.clearAllHighlights()
        
        if itm is not None:
            startLine = itm.data(0, self.StartLineRole)
            endLine = itm.data(0, self.EndLineRole)
            
            self.__editor.gotoLine(startLine, firstVisible=True,
                                   expand=True)
            self.__editor.setHighlight(startLine - 1, 0, endLine, -1)
    
    def __tryCompile(self, source, name):
        """
        Private method to attempt to compile the given source, first as an
        expression and then as a statement if the first approach fails.
        
        @param source source code string to be compiled
        @type str
        @param name name of the file containing the source
        @type str
        @return compiled code
        @rtype code object
        """
        try:
            c = compile(source, name, 'eval')
        except SyntaxError:
            c = compile(source, name, 'exec')
        return c
    
    def __disassembleObject(self, co, parentItem, parentName="", lasti=-1):
        """
        Private method to disassemble the given code object recursively.
        
        @param co code object to be disassembled
        @type code object
        @param parentItem reference to the parent item
        @type QTreeWidget or QTreeWidgetItem
        @param parentName name of the parent code object
        @type str
        @param lasti index of the instruction of a traceback
        @type int
        """
        if co.co_name == "<module>":
            title = os.path.basename(co.co_filename)
            name = ""
        else:
            if parentName:
                name = "{0}.{1}".format(parentName, co.co_name)
            else:
                name = co.co_name
            title = self.tr("Code Object '{0}'").format(name)
        titleItem = self.__createTitleItem(title, co.co_firstlineno,
                                           parentItem)
        lastStartItem = None
        for instr in dis.get_instructions(co):
            if instr.starts_line:
                if lastStartItem:
                    self.__updateItemEndLine(lastStartItem)
                lastStartItem = self.__createInstructionItem(
                    instr, titleItem, lasti=lasti)
            else:
                self.__createInstructionItem(instr, lastStartItem, lasti=lasti)
        if lastStartItem:
            self.__updateItemEndLine(lastStartItem)
        
        for x in co.co_consts:
            if hasattr(x, 'co_code'):
                self.__disassembleObject(x, titleItem, parentName=name,
                                         lasti=lasti)
        
        self.__updateItemEndLine(titleItem)
    
    @pyqtSlot()
    def preferencesChanged(self):
        """
        Public slot handling changes of the Disassembly viewer settings.
        """
        self.__errorColor = QBrush(
            Preferences.getPython("DisViewerErrorColor"))
        self.__currentInstructionColor = QBrush(
            Preferences.getPython("DisViewerCurrentColor"))
        self.__jumpTargetColor = QBrush(
            Preferences.getPython("DisViewerLabeledColor"))
        
        if self.isVisible():
            self.__loadDIS()
        
        self.__styleLabels()
    
    def __styleLabels(self):
        """
        Private method to style the info labels iaw. selected colors.
        """
        # current instruction
        self.__currentInfoLabel.setStyleSheet(
            "QLabel {{ color : {0}; }}".format(
                self.__currentInstructionColor.color().name()
            )
        )
        font = self.__currentInfoLabel.font()
        font.setItalic(True)
        self.__currentInfoLabel.setFont(font)
        
        # labeled instruction
        self.__labeledInfoLabel.setStyleSheet(
            "QLabel {{ color : {0}; }}".format(
                self.__jumpTargetColor.color().name()
            )
        )
        font = self.__labeledInfoLabel.font()
        font.setBold(True)
        self.__labeledInfoLabel.setFont(font)
    
    @pyqtSlot()
    def clear(self):
        """
        Public method to clear the display.
        """
        self.__disWidget.clear()