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
)


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.__infoLabel = QLabel(self.tr(
            "italic: current instruction\n"
            "bold: labelled instruction"
        ))
        self.__layout.addWidget(self.__infoLabel)
        
        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.__disWidget.itemClicked.connect(self.__disItemClicked)
        
        self.__vm.disViewerStateChanged.connect(self.__disViewerStateChanged)
        
        self.hide()
    
    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 __editorDoubleClicked(self, editor, pos, buttons):
        """
        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 pos position of the double click
        @type QPoint
        @param buttons mouse buttons that were double clicked
        @type Qt.MouseButtons
        """
        if editor is self.__editor and buttons == Qt.LeftButton:
            if editor.isModified():
                # reload the source
                QTimer.singleShot(0, self.__loadDIS)
            
            # highlight the corresponding entry
##            QTimer.singleShot(0, self.__selectItemForEditorSelection)
            QTimer.singleShot(0, self.__grabFocus)
    
    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.editorDoubleClickedEd.connect(self.__editorDoubleClicked)
            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.editorDoubleClickedEd.disconnect(
                self.__editorDoubleClicked)
            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 __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, QBrush(Qt.red))
        return itm
    
    def __createTitleItem(self, title, line):
        """
        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
        @return generated item
        @rtype QTreeWidgetItem
        """
        itm = QTreeWidgetItem(self.__disWidget, [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)
        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)
        
        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 not filename:
            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)
            
            self.__disassembleObject(codeObject, self.__disWidget, filename)
            QTimer.singleShot(0, self.__resizeColumns)
            
            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 __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, name="", 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 name name of the code object
        @type str
        @param lasti index of the instruction of a traceback
        @type int
        """
        if co.co_name == "<module>":
            title = (
                self.tr("Disassembly of module '{0}'")
                .format(os.path.basename(co.co_filename))
            )
        else:
            title = (
                self.tr("Disassembly of code object '{0}'")
                .format(co.co_name)
            )
        titleItem = self.__createTitleItem(title, co.co_firstlineno)
        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)
        
        self.__updateItemEndLine(titleItem)
        
        for x in co.co_consts:
            if hasattr(x, 'co_code'):
                self.__disassembleObject(x, self.__disWidget, lasti=lasti)