UI/PythonAstViewer.py

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

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

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

from __future__ import unicode_literals

try:
    str = unicode
except NameError:
    pass

import ast

from PyQt5.QtCore import Qt, QTimer
from PyQt5.QtGui import QCursor, QBrush
from PyQt5.QtWidgets import QTreeWidget, QApplication, QTreeWidgetItem, \
    QAbstractItemView, QWidget, QVBoxLayout

from ThirdParty.asttokens.asttokens import ASTTokens


# TODO: highlight code area in editor when a tree node is clicked
# TODO: jump to node when a double click in the editor is detected
#       (rebuild the tree, if the source code has changed)
class PythonAstViewer(QWidget):
    """
    Class implementing a widget to visualize the Python AST for some Python
    sources.
    """
    StartLineRole = Qt.UserRole
    StartIndexRole = Qt.UserRole + 1
    EndLineRole = Qt.UserRole + 2
    EndIndexRole = Qt.UserRole + 3
    
    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(PythonAstViewer, self).__init__(parent)
        
        self.__layout = QVBoxLayout(self)
        self.setLayout(self.__layout)
        self.__astWidget = QTreeWidget(self)
        self.__layout.addWidget(self.__astWidget)
        self.__layout.setContentsMargins(0, 0, 0, 0)
        
        self.__vm = viewmanager
        self.__vmConnected = False
        
        self.__editor = None
        self.__source = ""
        
        self.__astWidget.setHeaderLabels([self.tr("Node"),
                                          self.tr("Code Range")])
        self.__astWidget.setSortingEnabled(False)
        self.__astWidget.setSelectionBehavior(QAbstractItemView.SelectRows)
        self.__astWidget.setSelectionMode(QAbstractItemView.SingleSelection)
        self.__astWidget.setAlternatingRowColors(True)
        
        self.__vm.astViewerStateChanged.connect(self.__astViewerStateChanged)
        
        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:
            self.__editor = editor
            if self.__editor:
                self.__loadAST()
    
    def __editorSaved(self, editor):
        """
        Private slot to reload the AST 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.__loadAST()
    
    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 AST viewer.
        """
        super(PythonAstViewer, self).show()
        
        if not self.__vmConnected:
            self.__vm.editorChangedEd.connect(self.__editorChanged)
            self.__vm.editorSavedEd.connect(self.__editorSaved)
            self.__vmConnected = True
    
    def hide(self):
        """
        Public slot to hide the AST viewer.
        """
        super(PythonAstViewer, self).hide()
        
        if self.__vmConnected:
            self.__vm.editorChangedEd.disconnect(self.__editorChanged)
            self.__vm.editorSavedEd.disconnect(self.__editorSaved)
            self.__vmConnected = False
    
    def shutdown(self):
        """
        Public method to perform shutdown actions.
        """
        self.__editor = None
    
    def __astViewerStateChanged(self, on):
        """
        Private slot to toggle the display of the AST viewer.
        
        @param on flag indicating to show the AST
        @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.__loadAST()
        else:
            self.__editor = None
            self.hide()
    
    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.__astWidget, [error])
        itm.setFirstColumnSpanned(True)
        itm.setForeground(0, QBrush(Qt.red))
        return itm
    
    def __loadAST(self):
        """
        Private method to generate the AST from the source of the current
        editor and visualize it.
        """
        if not self.__editor:
            return
        
        self.__astWidget.clear()
        
        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
        
        QApplication.setOverrideCursor(QCursor(Qt.WaitCursor))
        try:
            # generate the AST
            root = ast.parse(source, self.__editor.getFileName(), "exec")
            self.__markTextRanges(root, source)
            astValid = True
        except Exception as exc:
            self.__createErrorItem(str(exc))
            astValid = False
        
        if astValid:
            self.setUpdatesEnabled(False)
            
            # populate the AST tree
            self.__populateNode(self.tr("Module"), root, self.__astWidget)
            self.__selectItemForEditorSelection()
            QTimer.singleShot(0, self.__resizeColumns)
            
            self.setUpdatesEnabled(True)
        
        QApplication.restoreOverrideCursor()
    
    def __populateNode(self, name, nodeOrFields, parent):
        """
        Private method to populate the tree view with a node.
        
        @param name name of the node
        @type str
        @param nodeOrFields reference to the node or a list node fields
        @type ast.AST or list
        @param parent reference to the parent item
        @type QTreeWidget or QTreeWidgetItem
        """
        if isinstance(nodeOrFields, ast.AST):
            fields = [(key, val) for key, val in ast.iter_fields(nodeOrFields)]
            value = nodeOrFields.__class__.__name__
        elif isinstance(nodeOrFields, list):
            fields = list(enumerate(nodeOrFields))
            if len(nodeOrFields) == 0:
                value = "[]"
            else:
                value = "[...]"
        else:
            fields = []
            value = repr(nodeOrFields)
        
        text = self.tr("{0}: {1}").format(name, value)
        itm = QTreeWidgetItem(parent, [text])
        itm.setExpanded(True)
        
        if hasattr(nodeOrFields, "lineno") and \
           hasattr(nodeOrFields, "col_offset"):
            itm.setData(0, self.StartLineRole, nodeOrFields.lineno)
            itm.setData(0, self.StartIndexRole, nodeOrFields.col_offset)
            startStr = self.tr("{0},{1}").format(
                nodeOrFields.lineno, nodeOrFields.col_offset)
            endStr = ""
            
            if hasattr(nodeOrFields, "end_lineno") and \
               hasattr(nodeOrFields, "end_col_offset"):
                itm.setData(0, self.EndLineRole, nodeOrFields.end_lineno)
                itm.setData(0, self.EndIndexRole,
                            nodeOrFields.end_col_offset)
                endStr = self.tr("{0},{1}").format(
                    nodeOrFields.end_lineno, nodeOrFields.end_col_offset)
            else:
                itm.setData(0, self.EndLineRole, nodeOrFields.lineno)
                itm.setData(0, self.EndIndexRole,
                            nodeOrFields.col_offset + 1)
            if endStr:
                rangeStr = self.tr("{0}  -  {1}").format(startStr, endStr)
            else:
                rangeStr = startStr
            
            itm.setText(1, rangeStr)
        
        for fieldName, fieldValue in fields:
            self.__populateNode(fieldName, fieldValue, itm)
    
    def __markTextRanges(self, tree, source):
        """
        Private method to modify the AST nodes with end_lineno and
        end_col_offset information.
        
        Note: The modifications are only done for nodes containing lineno and
        col_offset attributes.
        
        @param tree reference to the AST to be modified
        @type ast.AST
        @param source source code the AST was created from
        @type str
        """
        ASTTokens(source, tree=tree)
        for child in ast.walk(tree):
            if hasattr(child, 'last_token'):
                child.end_lineno, child.end_col_offset = child.last_token.end
                if hasattr(child, 'lineno'):
                    # Fixes problems with some nodes like binop
                    child.lineno, child.col_offset = child.first_token.start
    
    def __findClosestContainingNode(self, node, textRange):
        """
        Private method to search for the AST node that contains a range
        closest.
        
        @param node AST node to start searching at
        @type ast.AST
        @param textRange tuple giving the start and end positions
        @type tuple of (int, int, int, int)
        @return best matching node
        @rtype ast.AST
        """
        if textRange == (-1, -1, -1, -1):
            # no valid range, i.e. no selection
            return None
        
        # first look among children
        for child in ast.iter_child_nodes(node):
            result = self.__findClosestContainingNode(child, textRange)
            if result is not None:
                return result
        
        # no suitable child was found
        if hasattr(node, "lineno") and self.__rangeContainsSmaller(
            (node.lineno, node.col_offset, node.end_lineno,
             node.end_col_offset), textRange):
            return node
        else:
            # nope
            return None
    
    def __findClosestContainingItem(self, itm, textRange):
        """
        Private method to search for the tree item that contains a range
        closest.
        
        @param itm tree item to start searching at
        @type QTreeWidgetItem
        @param textRange tuple giving the start and end positions
        @type tuple of (int, int, int, int)
        @return best matching tree item
        @rtype QTreeWidgetItem
        """
        if textRange == (-1, -1, -1, -1):
            # no valid range, i.e. no selection
            return None
        
        # first look among children
        for index in range(itm.childCount()):
            child = itm.child(index)
            result = self.__findClosestContainingItem(child, textRange)
            if result is not None:
                return result
        
        # no suitable child was found
        lineno = itm.data(0, self.StartLineRole)
        if lineno is not None and self.__rangeContainsSmaller(
           (itm.data(0, self.StartLineRole), itm.data(0, self.StartIndexRole),
            itm.data(0, self.EndLineRole), itm.data(0, self.EndIndexRole)),
           textRange):
            return itm
        else:
            # nope
            return None
    
    def __resizeColumns(self):
        """
        Private method to resize the columns to suitable values.
        """
        for col in range(self.__astWidget.columnCount()):
            self.__astWidget.resizeColumnToContents(col)
        
        rangeSize = self.__astWidget.columnWidth(1) + 10
        # 10 px extra for the range
        nodeSize = max(400, self.__astWidget.viewport().width() - rangeSize)
        self.__astWidget.setColumnWidth(0, nodeSize)
        self.__astWidget.setColumnWidth(1, rangeSize)
    
    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 __rangeContainsSmaller(self, first, second):
        """
        Private method to check, if second is contained in first.
        
        @param first text range to check against
        @type tuple of (int, int, int, int)
        @param second text range to check for
        @type tuple of (int, int, int, int)
        @return flag indicating second is contained in first
        @rtype bool
        """
        firstStart = first[:2]
        firstEnd = first[2:]
        secondStart = second[:2]
        secondEnd = second[2:]

        return (
            (firstStart < secondStart and firstEnd > secondEnd) or
            (firstStart == secondStart and firstEnd > secondEnd) or
            (firstStart < secondStart and firstEnd == secondEnd)
        )
    
    def __clearSelection(self):
        """
        Private method to clear all selected items.
        """
        for itm in self.__astWidget.selectedItems():
            itm.setSelected(False)
    
    def __selectItemForEditorSelection(self):
        """
        Private slot to select the item corresponding to an editor selection.
        """
        # step 1: clear all selected items
        self.__clearSelection()
        
        # step 2: retrieve the editor selection
        selection = self.__editor.getSelection()
        # make the line numbers 1-based
        selection = (selection[0] + 1, selection[1],
                     selection[2] + 1, selection[3])
        
        # step 3: search the corresponding item, scroll to it and select it
        itm = self.__findClosestContainingItem(
            self.__astWidget.topLevelItem(0), selection)
        if itm:
            self.__astWidget.scrollToItem(
                itm, QAbstractItemView.PositionAtCenter)
            itm.setSelected(True)