# -*- coding: utf-8 -*-
"""
***************************************************************************
trees
---------------------
Date : November 2017
Copyright : (C) 2017 by Benjamin Jakimow
Email : benjamin.jakimow@geo.hu-berlin.de
***************************************************************************
* *
* This program is free software; you can redistribute it and/or modify *
* it under the terms of the GNU General Public License as published by *
* the Free Software Foundation; either version 2 of the License, or *
* (at your option) any later version. *
* *
***************************************************************************
"""
# noinspection PyPep8Naming
import os, pickle
from collections import OrderedDict
from qgis.core import *
from qgis.gui import *
from PyQt5.QtCore import *
from PyQt5.QtGui import *
from osgeo import gdal, osr
class TreeNode(QObject):
sigWillAddChildren = pyqtSignal(QObject, int, int)
sigAddedChildren = pyqtSignal(QObject, int, int)
sigWillRemoveChildren = pyqtSignal(QObject, int, int)
sigRemovedChildren = pyqtSignal(QObject, int, int)
sigUpdated = pyqtSignal(QObject)
def __init__(self, parentNode, name=None, values=None):
super(TreeNode, self).__init__()
self.mParent = parentNode
self.mChildren = []
self.mName = name
self.mValues = []
self.mIcon = None
self.mToolTip = None
if name:
self.setName(name)
if values:
self.setValues(values)
if isinstance(parentNode, TreeNode):
parentNode.appendChildNodes(self)
def nodeIndex(self):
return self.mParent.mChildren.index(self)
def next(self):
i = self.nodeIndex()
if i < len(self.mChildren.mChildren):
return self.mParent.mChildren[i + 1]
else:
return None
def previous(self):
i = self.nodeIndex()
if i > 0:
return self.mParent.mChildren[i - 1]
else:
return None
def detach(self):
"""
Detaches this TreeNode from its parent TreeNode
:return:
"""
if isinstance(self.mParent, TreeNode):
self.mParent.mChildren.remove(self)
self.setParentNode(None)
def appendChildNodes(self, listOfChildNodes):
self.insertChildNodes(len(self.mChildren), listOfChildNodes)
def insertChildNodes(self, index, listOfChildNodes):
assert index <= len(self.mChildren)
if isinstance(listOfChildNodes, TreeNode):
listOfChildNodes = [listOfChildNodes]
assert isinstance(listOfChildNodes, list)
l = len(listOfChildNodes)
idxLast = index + l - 1
self.sigWillAddChildren.emit(self, index, idxLast)
for i, node in enumerate(listOfChildNodes):
assert isinstance(node, TreeNode)
node.mParent = self
# connect node signals
node.sigWillAddChildren.connect(self.sigWillAddChildren)
node.sigAddedChildren.connect(self.sigAddedChildren)
node.sigWillRemoveChildren.connect(self.sigWillRemoveChildren)
node.sigRemovedChildren.connect(self.sigRemovedChildren)
node.sigUpdated.connect(self.sigUpdated)
self.mChildren.insert(index + i, node)
self.sigAddedChildren.emit(self, index, idxLast)
def removeChildNode(self, node):
assert node in self.mChildren
i = self.mChildren.index(node)
self.removeChildNodes(i, 1)
def removeChildNodes(self, row, count):
if row < 0 or count <= 0:
return False
rowLast = row + count - 1
if rowLast >= self.childCount():
return False
self.sigWillRemoveChildren.emit(self, row, rowLast)
to_remove = self.childNodes()[row:rowLast + 1]
for n in to_remove:
self.mChildren.remove(n)
# n.mParent = None
self.sigRemovedChildren.emit(self, row, rowLast)
def setToolTip(self, toolTip):
self.mToolTip = toolTip
def toolTip(self):
return self.mToolTip
def parentNode(self):
return self.mParent
def setParentNode(self, treeNode):
assert isinstance(treeNode, TreeNode)
self.mParent = treeNode
def setIcon(self, icon):
self.mIcon = icon
def icon(self):
return self.mIcon
def setName(self, name):
self.mName = name
def name(self):
return self.mName
def contextMenu(self):
return None
def setValues(self, listOfValues):
if not isinstance(listOfValues, list):
listOfValues = [listOfValues]
self.mValues = listOfValues[:]
def values(self):
return self.mValues[:]
def childCount(self):
return len(self.mChildren)
def childNodes(self):
return self.mChildren[:]
def findChildNodes(self, type, recursive=True):
results = []
for node in self.mChildren:
if isinstance(node, type):
results.append(node)
if recursive:
results.extend(node.findChildNodes(type, recursive=True))
return results
class TreeModel(QAbstractItemModel):
def __init__(self, parent=None, rootNode=None):
super(TreeModel, self).__init__(parent)
self.mColumnNames = ['Node', 'Value']
self.mRootNode = rootNode if isinstance(rootNode, TreeNode) else TreeNode(None)
self.mRootNode.sigWillAddChildren.connect(self.nodeWillAddChildren)
self.mRootNode.sigAddedChildren.connect(self.nodeAddedChildren)
self.mRootNode.sigWillRemoveChildren.connect(self.nodeWillRemoveChildren)
self.mRootNode.sigRemovedChildren.connect(self.nodeRemovedChildren)
self.mRootNode.sigUpdated.connect(self.nodeUpdated)
self.mTreeView = None
if isinstance(parent, QTreeView):
self.connectTreeView(parent)
def nodeWillAddChildren(self, node, idx1, idxL):
idxNode = self.node2idx(node)
self.beginInsertRows(idxNode, idx1, idxL)
def nodeAddedChildren(self, node, idx1, idxL):
self.endInsertRows()
# for i in range(idx1, idxL+1):
for n in node.childNodes():
# self.setColumnSpan(node)
pass
def nodeWillRemoveChildren(self, node, idx1, idxL):
idxNode = self.node2idx(node)
self.beginRemoveRows(idxNode, idx1, idxL)
def nodeRemovedChildren(self, node, idx1, idxL):
self.endRemoveRows()
def nodeUpdated(self, node):
idxNode = self.node2idx(node)
self.dataChanged.emit(idxNode, idxNode)
self.setColumnSpan(node)
def headerData(self, section, orientation, role):
if orientation == Qt.Horizontal and role == Qt.DisplayRole:
if len(self.mColumnNames) > section:
return self.mColumnNames[section]
else:
return ''
else:
return None
def parent(self, index):
if not index.isValid():
return QModelIndex()
node = self.idx2node(index)
if not isinstance(node, TreeNode):
return QModelIndex()
parentNode = node.parentNode()
if not isinstance(parentNode, TreeNode):
return QModelIndex()
return self.node2idx(parentNode)
if node not in parentNode.mChildren:
return QModelIndex
row = parentNode.mChildren.index(node)
return self.createIndex(row, 0, parentNode)
def rowCount(self, index):
node = self.idx2node(index)
return len(node.mChildren) if isinstance(node, TreeNode) else 0
def hasChildren(self, index):
node = self.idx2node(index)
return isinstance(node, TreeNode) and len(node.mChildren) > 0
def columnNames(self):
return self.mColumnNames
def columnCount(self, index):
return len(self.mColumnNames)
def connectTreeView(self, treeView):
self.mTreeView = treeView
def setColumnSpan(self, node, span=None):
if isinstance(self.mTreeView, QTreeView) \
and isinstance(node, TreeNode) \
and isinstance(node.parentNode(), TreeNode):
idxNode = self.node2idx(node)
idxParent = self.node2idx(node.parentNode())
if not isinstance(span, bool):
span = len(node.values()) == 0
self.mTreeView.setFirstColumnSpanned(idxNode.row(), idxParent, span)
for n in node.childNodes():
self.setColumnSpan(n)
def index(self, row, column, parentIndex=None):
if parentIndex is None:
parentNode = self.mRootNode
else:
parentNode = self.idx2node(parentIndex)
if row < 0 or row >= parentNode.childCount():
return QModelIndex()
if column < 0 or column >= len(self.mColumnNames):
return QModelIndex()
if isinstance(parentNode, TreeNode) and row < len(parentNode.mChildren):
return self.createIndex(row, column, parentNode.mChildren[row])
else:
return QModelIndex()
def findParentNode(self, node, parentNodeType):
assert isinstance(node, TreeNode)
while True:
if isinstance(node, parentNodeType):
return node
if not isinstance(node.parentNode(), TreeNode):
return None
node = node.parentNode()
def indexes2nodes(self, indexes):
assert isinstance(indexes, list)
nodes = []
for idx in indexes:
n = self.idx2node(idx)
if n not in nodes:
nodes.append(n)
return nodes
def expandNode(self, node, expand=True, recursive=True):
assert isinstance(node, TreeNode)
if isinstance(self.mTreeView, QTreeView):
idx = self.node2idx(node)
self.mTreeView.setExpanded(idx, expand)
if recursive:
for n in node.childNodes():
self.expandNode(n, expand=expand, recursive=recursive)
def nodes2indexes(self, nodes):
return [self.node2idx(n) for n in nodes]
def idx2node(self, index):
if not index.isValid():
return self.mRootNode
else:
return index.internalPointer()
def node2idx(self, node):
assert isinstance(node, TreeNode)
if node == self.mRootNode:
return QModelIndex()
else:
parentNode = node.parentNode()
assert isinstance(parentNode, TreeNode)
if node not in parentNode.mChildren:
return QModelIndex()
r = parentNode.mChildren.index(node)
return self.createIndex(r, 0, node)
def data(self, index, role):
node = self.idx2node(index)
col = index.column()
if role == Qt.UserRole:
return node
if col == 0:
if role in [Qt.DisplayRole, Qt.EditRole]:
return node.name()
if role == Qt.DecorationRole:
return node.icon()
if role == Qt.ToolTipRole:
return node.toolTip()
if col > 0:
i = col - 1
if role in [Qt.DisplayRole, Qt.EditRole] and len(node.values()) > i:
return str(node.values()[i])
def flags(self, index):
if not index.isValid():
return Qt.NoItemFlags
node = self.idx2node(index)
return Qt.ItemIsEnabled | Qt.ItemIsSelectable
class TreeView(QTreeView):
def __init__(self, *args, **kwds):
super(TreeView, self).__init__(*args, **kwds)