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

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 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, parent=None):
        """
        Constructor
        
        @param viewmanager reference to the viewmanager object
        @type ViewManager
        @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.__vm = viewmanager
        self.__vmConnected = False
        
        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)
        
        self.__vm.disViewerStateChanged.connect(self.__disViewerStateChanged)
        
        self.hide()
    
    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 __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 not self.__vmConnected:
            self.__vm.editorChangedEd.connect(self.__editorChanged)
            self.__vm.editorSavedEd.connect(self.__editorSaved)
            self.__vm.editorLineChangedEd.connect(self.__editorLineChanged)
            self.__vmConnected = True
    
    def hide(self):
        """
        Public slot to hide the DIS viewer.
        """
        super(PythonDisViewer, self).hide()
        
        if self.__editor:
            self.__editor.clearAllHighlights()
        
        if self.__vmConnected:
            self.__vm.editorChangedEd.disconnect(self.__editorChanged)
            self.__vm.editorSavedEd.disconnect(self.__editorSaved)
            self.__vm.editorLineChangedEd.disconnect(self.__editorLineChanged)
            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
        """
        editor = self.__vm.activeWindow()
        if on and editor and editor.isPyFile():
            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 not self.__editor:
            return
        
        self.__disWidget.clear()
        self.__editor.clearAllHighlights()
        
        if not self.__editor.isPyFile():
            self.__createErrorItem(self.tr(
                "The current editor text does not contain Python source."
            ))
            return
        
        source = self.__editor.text()
        if not source.strip():
            # empty editor or white space only
            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()
        
        self.__grabFocus()
    
    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)
    
    def __grabFocus(self):
        """
        Private method to grab the input focus.
        """
        self.__disWidget.setFocus(Qt.OtherFocusReason)
    
    @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
        """
        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()