QtForum.org

Source code

# -*- coding: utf-8 -*-
"""
Created on Wed Jun 27 20:01:55 2012

@author: dpwrussell
"""
from PyQt4 import QtCore, QtGui
import sys


class TreeItem(object):
    def __init__(self, data, parent = None):
        
        self.childItems = []
        self.itemData = data
        self.parentItem = parent

    #Append a child TreeItem to this item
    def appendChild(self, child):
        self.childItems.append(child)
    
    #Get the child at the specified row number
    def child(self, row):
        return self.childItems[row]

    #Get the number of children    
    def childCount(self):
        return len(self.childItems)

    #Get the number of columns for the item    
    def columnCount(self):
        return len(self.itemData)
    
    #Get the data (QVariant) for the specified column of this item
    def data(self, column):
        return self.itemData[column]
    
    #Get the row number of this item (within its parent)
    def row(self):
        # IF it's not the root item
        if (self.parentItem):
            #Return the index of this item
            return self.parentItem.childItems.index(self)
        #Apparently this zero is never used by the model so it's safe despite being the same as the first child
        return 0

    #Get the parent (TreeItem)    
    def parent(self):
        return self.parentItem

class TreeModel(QtCore.QAbstractItemModel):

    def __init__(self, data, parent = None):
        QtCore.QAbstractItemModel.__init__(self, parent)
        #Root Item is the false top-level item, also contains the header information
        self.rootItem = TreeItem([QtCore.QVariant("Name"),QtCore.QVariant("ID")])
        
        self.__setupModelData([], self.rootItem)
        
        print self.rootItem.childCount()
    
    #Get the data (QVariant) for the specified index (QModelIndex) and role
    def data(self, index, role):
        #If the index is not valid
        if not index.isValid():
            #Return an invlalid QVariant
            return QtCore.QVariant()
        
        if role != QtCore.Qt.DisplayRole:
            return QtCore.QVariant()
        
        #Get the actual item using the index
        item = index.internalPointer()
        #Return the data for the specified column
        return item.data(index.column())
    
    #Get flags for the specified index (QModelIndex)
    def flags(self, index):
        #If the index is not valid
        if not index.isValid():
            #Return an invlalid QVariant
            return 0
            
        return QtCore.Qt.ItemIsEnabled | QtCore.Qt.ItemIsSelectable
    
    #Get the headerData (QVariant) for the specified section (called section as I guess for different oreintations this is a row or column), orientation and role
    #TODO Temporarily removed default from role (role = QtCore.Qt.DisplayRole())
    def headerData(self, section, orientation, role):
        if orientation == QtCore.Qt.Horizontal and role == QtCore.Qt.DisplayRole:
            return self.rootItem.data(section)
        return QtCore.QVariant()
    
    #Get index (QModelIndex) of the specified row, column and parent (QmodelIndex)
    def index(self, row, column, parent=QtCore.QModelIndex()):
        #If the model does not have this index
        if not self.hasIndex(row, column, parent):
            return QtCore.QModelIndex()
        
        #First, find the parent
        #If the parent is not valid
        if not parent.isValid():
            #Root is assumed to be the false top-level item
            parentItem = self.rootItem
        else:
            #
            parentItem = parent.internalPointer()
            
        #Second, get the child at the specified row
        childItem = parentItem.child(row)
        
        #Finally, ensure that there was a child
        if childItem:
            #Return a newly created index for that child
            return self.createIndex(row, column, childItem)
        else:
            return QtCore.QModelIndex()

    #Get index of parent (QModelIndex) of index
    def parent(self, index=QtCore.QModelIndex()):
        if not index.isValid():
            return QtCore.QModelIndex()
        
        childItem = index.internalPointer()
        parentItem = childItem.parent()
        
        if parentItem == self.rootItem:
            return QtCore.QModelIndex()
        
        return self.createIndex(parentItem.row(), 0, parentItem)
    
    #Get the number of rows for parent index
    def rowCount(self, parent=QtCore.QModelIndex()):
        #By convention, only column zero has children
        if parent.column > 0:
            return 0

        #If the parent is not valid        
        if not parent.isValid():
            #Set the parent as the rootItem as we want to return the number of top-level items
            parentItem = self.rootItem
        else:
            parentItem = parent.internalPointer()
        
        return parentItem.childCount()
    
    #get the number of columns for parent at index?
    def columnCount(self, parent=QtCore.QModelIndex()):
        #If the parent is not valid  
        if not parent.isValid():
            if parent.internalPointer():
                print "At least it had an internal pointter"
                
            else:
                print "Didn't have an internal pointer"
            #Return the number of columns in the rootItem, I guess this matches up with what rowCount does with the same problem
            return self.rootItem.columnCount()
        else:
            #Return number of columns in the parent
            return parent.internalPointer().columnCount()

    #Creates some dummy data, currently 8 tag sets (0-7) then 2 (0-1)tags as children of those.
    #Tags are numbered with the tagset number then the tag number which is visually unique for easy testing
    #Tagsets are added to the rootItem (parent below)
    # The data parameter does nothing at the moment, as the data is hardcoded    
    def __setupModelData(self, data, parent):        
        
        #5 Tagsets
        for i in range(8):
            tscols = []
            tscols.append(QtCore.QVariant("TagSet" + str(i)))
            tscols.append(QtCore.QVariant("IDTagSet" + str(i)))
        
            tsitem = TreeItem(tscols, parent)
            parent.appendChild(tsitem)
        
            #2 Tags in each
            for j in range(2):
                tcols = []
                tcols.append(QtCore.QVariant("Tag" + str(i) + "-" + str(j)))
                tcols.append(QtCore.QVariant("IDTag" + str(i) + "-" +  str(j)))
                
                titem = TreeItem(tcols, tsitem)
                tsitem.appendChild(titem)
        #Make sure the tree is populated by print to terminal (Seems OK)
        printTree(parent)
        
def printTree(item, tab=""):
    print tab,item.data(0).toPyObject()
    children = item.childItems
    for c in children:
        printTree(c, tab + "\t")
        
        
if __name__ == "__main__":

    app = QtGui.QApplication(sys.argv)
    model = TreeModel([])
    
    #initially verify the tree works by using a predefined model
#    dirsModel = QtGui.QFileSystemModel()
#    dirsModel.setFilter(QtCore.QDir.AllDirs | QtCore.QDir.Dirs | QtCore.QDir.Drives | QtCore.QDir.NoDotAndDotDot)
#    dirsModel.setRootPath("/")
    
    tree = QtGui.QTreeView()
    tree.setModel(model)
    tree.setAnimated(False)
    tree.setIndentation(20)
    tree.setSortingEnabled(True)

    tree.setWindowTitle("Test")
    tree.resize(640, 480)
    tree.show()
    sys.exit(app.exec_())