from PyQt4.QtCore import *
from PyQt4.QtGui import *
from c3 import Visitor, astnodes

class AstModelBuilder:
   def __init__(self):
      self.functionIco = QIcon(QPixmap('functionicon.png').scaled(32, 32))
      self.variableIco = QIcon(QPixmap('variableicon.png').scaled(32, 32))
      self.visitor = Visitor(self.p1, self.p2)
      self.model = QStandardItemModel()
      self.model.setHorizontalHeaderLabels(['Object', 'Type'])
   def build(self, pkg):
      #self.model.clear()
      c = self.model.rowCount()
      self.model.removeRows(0, c)
      self.curItem = self.model.invisibleRootItem()
      self.visitor.visit(pkg)
   def p1(self, node):
      if type(node) is astnodes.Variable:
         i = QStandardItem(self.variableIco, str(node))
      elif type(node) is astnodes.Function:
         i = QStandardItem(self.functionIco, str(node))
      elif type(node) is astnodes.Package:
         i = QStandardItem(str(node))
      else:
         return
      typ = str(node.typ) if hasattr(node, 'typ') else ''
      ti = QStandardItem(str(typ))
      ti.setData(node)
      i.setData(node)
      self.curItem.appendRow([i, ti])
      self.curItem = i
   def p2(self, node):
      if type(node) in [astnodes.Variable, astnodes.Function, astnodes.Package]:
         self.curItem = self.curItem.parent()

# The actual widget:
class AstViewer(QTreeView):
   sigNodeSelected = pyqtSignal(object)
   def __init__(self, parent=None):
      super(AstViewer, self).__init__(parent)
      self.clicked.connect(self.selectHandler)
      self.modelBuilder = AstModelBuilder()
      self.setModel(self.modelBuilder.model)

   def setAst(self, ast):
      """ Create a new model and add all ast elements to it """
      self.modelBuilder.build(ast)
      self.expandAll()

   def selectHandler(self, index):
      if not index.isValid():
         return
      model = self.model()
      item = model.itemFromIndex(index)
      node = item.data()
      self.sigNodeSelected.emit(node)