tree.py

tree.py

import sys
from PyQt5 import QtCore, QtGui, QtWidgets


class Node(object):
    def __init__(self, name, parent=None):
        self._name = name
        self._parent = parent
        self._children = []

        # Add this node to it's given parent
        if parent is not None:
            parent.addChild(self)

    def typeInfo(self):
        return "NODE"

    def addChild(self, child):
        self._children.append(child)

    def insertChild(self, position, child):
        if position < 0 or position > len(self._children):
            return False

        self._children.insert(position, child)
        child._parent = self
        return True

    def removeChild(self, position):
        if position < 0 or position > len(self._children):
            return False

        child = self._children.pop(position)
        child._parent = None
        return True

    def name(self):
        return self._name

    def setName(self, name):
        self._name = name

    def child(self, row):
        return self._children[row]

    def childCount(self):
        return len(self._children)

    def parent(self):
        return self._parent

    def row(self):
        if self._parent is not None:
            return self._parent._children.index(self)

    def log(self, tabLevel=-1):
        output = ""
        tabLevel += 1

        for i in range(tabLevel):
            output += "\t"

        output += "|---" + self._name + "\n"
        for child in self._children:
            output += child.log(tabLevel)

        tabLevel -= 1
        return output

    def __repr__(self):
        return self.log()


class TransformNode(Node):
    def __init__(self, name, parent=None):
        super(TransformNode, self).__init__(name, parent)

    def typeInfo(self):
        return "Transform"


class CameraNode(Node):
    def __init__(self, name, parent=None):
        super(CameraNode, self).__init__(name, parent)

    def typeInfo(self):
        return "Camera"


class LightNode(Node):
    def __init__(self, name, parent=None):
        super(LightNode, self).__init__(name, parent)

    def typeInfo(self):
        return "Light"


class TreeModel(QtCore.QAbstractItemModel):
    def __init__(self, root, parent=None):
        super(TreeModel, self).__init__(parent)
        self._root_node = root

    def parent(self, index):
        node = self.getNode(index)
        parent_node = node.parent()

        if parent_node == self._root_node:
            return QtCore.QModelIndex()

        return self.createIndex(parent_node.row(), 0, parent_node)

    def index(self, row, column, parent):
        parent_node = self.getNode(parent)
        child_item = parent_node.child(row)

        if child_item:
            return self.createIndex(row, column, child_item)
        else:
            return QtCore.QModelIndex()

    def getNode(self, index):
        if index.isValid():
            node = index.internalPointer()
            if node:
                return node
        return self._root_node

    def insertRows(self, position, rows, parent=QtCore.QModelIndex()):
        parent_node = self.getNode(parent)
        child_count = parent_node.childCount()

        self.beginInsertRows(parent, position, position + rows - 1)

        for row in range(rows):
            child_node = Node("Untitled" + str(child_count))
            # Should this be verified?
            success = parent_node.insertChild(position, child_node)

        self.endInsertRows()
        return success

    def insertLights(self, position, rows, parent=QtCore.QModelIndex()):
        parent_node = self.getNode(parent)
        child_count = parent_node.childCount()

        self.beginInsertRows(parent, position, position + rows - 1)

        for row in range(rows):
            child_node = LightNode("Light" + str(child_count))
            # Should this be verified?
            success = parent_node.insertChild(position, child_node)

        self.endInsertRows()
        return success

    def removeRows(self, position, rows, parent=QtCore.QModelIndex()):
        parent_node = self.getNode(parent)
        self.beginRemoveRows(parent, position, position + rows - 1)

        for row in range(rows):
            # Should this be verified?
            success = parent_node.removeChild(position)

        self.endRemoveRows()
        return success


    def rowCount(self, parent=None):
        if not parent.isValid():
            parent_node = self._root_node
        else:
            parent_node = parent.internalPointer()

        return parent_node.childCount()

    def columnCount(self, parent=None):
        return 2

    def flags(self, index):
        return QtCore.Qt.ItemIsEnabled | QtCore.Qt.ItemIsSelectable | QtCore.Qt.ItemIsEditable

    def headerData(self, section, orientation, role=None):
        if role == QtCore.Qt.DisplayRole:
            if section == 0:
                return 'Scenegraph'
            else:
                return 'TypeInfo'

    def data(self, index, role):
        if not index.isValid():
            return None

        node = index.internalPointer()

        if role == QtCore.Qt.DisplayRole or role == QtCore.Qt.EditRole:
            if index.column() == 0:
                return node.name()
            else:
                # Fill in extra columns with extra data.
                # return the type of the node. Can fill this with data later.
                return node.typeInfo()

        # Decoration, icon classes. Can get that when you have types of
        # assets loading and stuff.
        # if role == QtCore.Qt.DecorationRole:
        #     if index.column() == 0:
        #         type_info = node.typeInfo()
        #
        #         if type_info == 'Light':
        #             return QtGui.QIcon(QtGui.QPixmap(":/Light.png"))
        #
        #         if type_info == 'Transform':
        #             return QtGui.QIcon(QtGui.QPixmap(":/Transform.png"))
        #
        #         if type_info == 'Camera':
        #             return QtGui.QIcon(QtGui.QPixmap(":/Camera.png"))

    def setData(self, index, value, role=QtCore.Qt.EditRole):
        if index.isValid():
            if role == QtCore.Qt.EditRole:
                node = index.internalPointer()
                node.setName(value)
                return True
        return False



if __name__ == "__main__":

    app = QtWidgets.QApplication(sys.argv)

    app.setStyle('plastique')

    # Put in a dialog.
    # Then put in a qlineedit.
    # Do sorting & filtering.


    rootNode = Node("Hips")
    childNode0 = Node("RightPirateLeg", rootNode)
    childNode1 = Node("RightFoot", childNode0)

    childNode2 = TransformNode("LeftFemur", rootNode)
    childNode3 = Node("LeftTibia", childNode2)
    childNode4 = CameraNode("LeftFoot", childNode3)
    childNode5 = LightNode("LeftFootEnd", childNode4)
    proxyModel = QtCore.QSortFilterProxyModel()

    model = TreeModel(rootNode)
    proxyModel.setSourceModel(model)
    proxyModel.setDynamicSortFilter(True)
    proxyModel.setFilterCaseSensitivity(QtCore.Qt.CaseInsensitive)

    print(rootNode)

    tree_view = QtWidgets.QTreeView()
    tree_view.show()
    tree_view.setModel(model)

    model.insertRows(1, 5)
    model.insertLights(1, 5)

    model.removeRows(1, 5)

    right_pirate_leg = model.index(0, 0, QtCore.QModelIndex())
    model.insertRows(1, 5, right_pirate_leg)
    model.insertLights(1, 5, right_pirate_leg)

    sys.exit(app.exec_())