ex.py

#!/usr/bin/env python3

#TODO: """DocString if there is one"""

import string
import functools
import random
import sys

from PyQt5.QtCore import (Qt, QAbstractItemModel, QVariant,
                          QModelIndex, pyqtSignal, pyqtSlot, QSize)
from PyQt5.QtWidgets import (QWidget, QApplication, QTreeView, QMainWindow,
                             QHBoxLayout, QVBoxLayout, QPushButton)

# See file COPYING in this source tree
__copyright__ = 'Copyright 2016, EPC Power Corp.'
__license__ = 'GPLv2+'


class TreeNode:
    def __init__(self,  tx=False, parent=None):
        self.last = None

        self.tx = tx

        self.tree_parent = None
        self.set_parent(parent)
        self.children = []

    def set_parent(self, parent):
        self.tree_parent = parent
        if self.tree_parent is not None:
            self.tree_parent.append_child(self)

    def append_child(self, child):
        self.children.append(child)
        child.tree_parent = self

    def child_at_row(self, row):
        try:
            return self.children[row]
        except IndexError:
            return None

    def row_of_child(self, child):
        for i, item in enumerate(self.children):
            if item == child:
                return i
        return -1

    def remove_child(self, row=None, child=None):
        if child is None:
            child = self.children[row]

        self.children.remove(child)

        return True

    def traverse(self, call_this, payload=None):
        for child in self.children:
            if len(child.children) == 0:
                call_this(child, payload)
            else:
                child.traverse(child, call_this, payload)

    def __len__(self):
        return len(self.children)


unique_role = Qt.UserRole


class PyQAbstractItemModel(QAbstractItemModel):
    root_changed = pyqtSignal(TreeNode)

    def __init__(self, root, checkbox_columns=None, editable_columns=None,
                 alignment=None, parent=None):
        QAbstractItemModel.__init__(self, parent=parent)

        self.root = root
        self.checkbox_columns = checkbox_columns
        self.editable_columns = editable_columns

        if alignment is not None:
            self.alignment = alignment
        else:
            self.alignment = Qt.AlignTop | Qt.AlignLeft

        self.index_from_node_cache = {}

        self.role_functions = {
            Qt.DisplayRole: self.data_display,
            unique_role: self.data_unique,
            Qt.TextAlignmentRole: lambda index: int(self.alignment),
            Qt.CheckStateRole: self.data_check_state,
            Qt.EditRole: self.data_edit,
            Qt.SizeHintRole: self.data_size_hint
        }

    def headerData(self, section, orientation, role):
        if orientation == Qt.Horizontal and role == Qt.DisplayRole:
            return QVariant(self.headers[section])
        return QVariant()

    def data_display(self, index):
        node = index.internalPointer()

        try:
            return node.fields[index.column()]
        except IndexError:
            return None

    def data_unique(self, index):
        return index.internalPointer().unique()

    def data_check_state(self, index):
        if self.checkbox_columns is not None:
            if self.checkbox_columns[index.column()]:
                node = index.internalPointer()
                try:
                    return node.checked(index.column())
                except AttributeError:
                    return None

    def data_edit(self, index):
        node = index.internalPointer()
        try:
            get = node.get_human_value
        except AttributeError:
            value = node.fields[index.column()]
        else:
            try:
                value = get()
            except TypeError:
                value = ''

        if value is None:
            value = ''
        else:
            value = str(value)

        return value

    def data_size_hint(self, index):
        return QSize(50, 50)

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

        try:
            return self.role_functions[role](index=index)
        except KeyError:
            return None

    def flags(self, index):
        flags = QAbstractItemModel.flags(self, index)

        if not index.isValid():
            return flags

        if self.editable_columns is not None:
            if self.editable_columns[index.column()]:
                flags |= Qt.ItemIsEditable

        if self.checkbox_columns is not None:
            if self.checkbox_columns[index.column()]:
                flags |= Qt.ItemIsUserCheckable

        return flags

    def index(self, row, column, parent):
        # TODO: commented out stuff ought to be good rather than
        #       breaking stuff.
        #
        #       http://stackoverflow.com/questions/26680168/pyqt-treeview-index-error-removing-last-row

        if not self.hasIndex(row, column, parent):
            return QModelIndex()

        # if not parent.isValid():
        #     return QModelIndex()

        # if row < 0 or column < 0:
        #     return QModelIndex()

        node = self.node_from_index(parent)
        child = node.child_at_row(row)

        if child is None:
            return QModelIndex()

        return self.createIndex(row, column, child)

    def columnCount(self, parent):
        return len(self.headers)

    def rowCount(self, parent):
        # TODO: this seems pretty particular to my present model
        #       "the second column should NOT have the same children
        #       as the first column in a row"
        #       https://github.com/bgr/PyQt5_modeltest/blob/62bc86edbad065097c4835ceb4eee5fa3754f527/modeltest.py#L222
        #
        #       then again, the Qt example does just this
        #       http://doc.qt.io/qt-5/qtwidgets-itemviews-simpletreemodel-example.html
        if parent.column() > 0:
            return 0

        node = self.node_from_index(parent)
        if node is None:
            return 0
        return len(node)

    def parent(self, child):
        if not child.isValid():
            return QModelIndex()

        node = self.node_from_index(child)

        if node is None:
            return QModelIndex()

        parent = node.tree_parent

        if parent in [None, self.root]:
            return QModelIndex()

        grandparent = parent.tree_parent
        if grandparent is None:
            return QModelIndex()
        row = grandparent.row_of_child(parent)

        assert row != - 1
        return self.createIndex(row, 0, parent)

    def node_from_index(self, index):
        if index.isValid():
            return index.internalPointer()
        else:
            return self.root

    def index_from_node(self, node):
        # TODO  make up another role for identification?
        try:
            index = self.index_from_node_cache[node]
        except KeyError:
            if node is self.root:
                index = QModelIndex()
            else:
                index = self.match(self.index(0, 0, QModelIndex()),
                                   unique_role,
                                   node.unique(),
                                   1,
                                   Qt.MatchRecursive)[0]

                self.index_from_node_cache[node] = index

        return index

    @pyqtSlot(TreeNode, int, TreeNode, int, list)
    def changed(self, start_node, start_column, end_node, end_column, roles):
        start_index = self.index_from_node(start_node)
        start_row = start_index.row()
        start_parent = start_index.parent()
        start_index = self.index(start_row, start_column, start_parent)

        if end_node is start_node:
            end_row = start_row
            end_parent = start_parent
        else:
            end_index = self.index_from_node(end_node)
            end_row = end_index.row()
            end_parent = end_index.parent()

        end_index = self.index(end_row, end_column, end_parent)

        self.dataChanged.emit(start_index, end_index, roles)
        print('dataChanged() emitted')
        for name, index in [('Start', start_index), ('End', end_index)]:
            print('{:5s}:r{},c{} parent({})'.format(name,
                                                    index.row(),
                                                    index.column(),
                                                    index.parent()))
        print('Parents equal: {}'.format(start_index.parent()
                                         == end_index.parent()))

    @pyqtSlot(TreeNode, int, int)
    def begin_insert_rows(self, parent, start_row, end_row):
        self.beginInsertRows(self.index_from_node(parent), start_row, end_row)

    @pyqtSlot()
    def end_insert_rows(self):
        self.index_from_node_cache = {}
        self.endInsertRows()

    @pyqtSlot(TreeNode, int, int)
    def begin_remove_rows(self, parent, start_row, end_row):
        self.beginRemoveRows(self.index_from_node(parent), start_row, end_row)

    @pyqtSlot()
    def end_remove_rows(self):
        self.index_from_node_cache = {}
        self.endRemoveRows()

    @pyqtSlot()
    def set_root(self, root):
        self.beginResetModel()
        self.root = root
        self.endResetModel()
        self.root_changed.emit(root)


class AbstractColumns:
    def __init__(self, **kwargs):
        for member in self._members:
            try:
                value = kwargs[member]
            except KeyError:
                value = None
            finally:
                setattr(self, member, value)

        object.__setattr__(self, '_length', len(self.__dict__))

        invalid_parameters = set(kwargs.keys()) - set(self.__dict__.keys())
        if len(invalid_parameters):
            raise ValueError('Invalid parameter{} passed: {}'.format(
                's' if len(invalid_parameters) > 1 else '',
                ', '.join(invalid_parameters)))

    @classmethod
    def __len__(cls):
        return len(cls._members)

    @classmethod
    def indexes(cls):
        return cls(**dict(zip(cls._members, range(len(cls._members)))))

    @classmethod
    def fill(cls, value):
        return cls(**dict(zip(cls._members, [value] * len(cls._members))))

    def __iter__(self):
        for i in range(len(self)):
            yield self[i]

    @functools.lru_cache(maxsize=None)
    def index_from_attribute(self, index):
        for attribute in self.__class__._members:
            if index == getattr(self.__class__.indexes, attribute):
                return attribute

        raise IndexError('column index out of range')

    def __getitem__(self, index):
        if index < 0:
            index += len(self)
        return getattr(self, self.index_from_attribute(index))

    def __setitem__(self, index, value):
        if index < 0:
            index += len(self)
        return setattr(self, self.index_from_attribute(index), value)

    def __getattr__(self, name, value):
        if name in self._members:
            object.__getattr__(self, name, value)
        else:
            raise TypeError("Attempted to get attribute {}"
                            .format(name))

    def __setattr__(self, name, value):
        if name in self._members:
            object.__setattr__(self, name, value)
        else:
            raise TypeError("Attempted to set attribute {}"
                            .format(name))


class Columns(AbstractColumns):
    _members = ['name', 'letter', 'number']

Columns.indexes = Columns.indexes()


class Node(TreeNode):
    def __init__(self, name, letter='a', number=0):
        TreeNode.__init__(self)
        self.fields = Columns(name=name, letter=letter, number=number)

        self.possibilities = Columns(letter=string.ascii_letters.lower(),
                                     number=range(10))

    def randomize(self, _):
        letters = set(self.possibilities.letter)
        letters.remove(self.fields.letter)
        self.fields.letter = random.choice(list(letters))

        numbers = set(self.possibilities.number)
        numbers.remove(self.fields.number)
        self.fields.number = random.choice(list(numbers))

    def unique(self):
        return object.__repr__(self)


class Model(PyQAbstractItemModel):
    def __init__(self, root, parent=None):
        PyQAbstractItemModel.__init__(self, root=root, parent=parent)

        self.headers = Columns(name='Name', letter='Letter', number='Number')

    def data_clicked(self):
        self.root.traverse(call_this=Node.randomize)

    def single_clicked(self):
        node = self.root.children[1]
        self.changed(node, 1, node, 1, [])

    def double_clicked(self):
        node = self.root.children[1]
        self.changed(node, 1, node, 2, [])


def main():
    app = QApplication(sys.argv)

    root = Node(name='Root', letter='-', number=-1)
    a = Node(name='A', letter='a', number=1)
    b = Node(name='B', letter='b', number=2)
    c = Node(name='C', letter='c', number=3)
    root.append_child(a)
    root.append_child(b)
    root.append_child(c)

    model = Model(root=root)

    window = QWidget()

    vlayout = QVBoxLayout()
    window.setLayout(vlayout)

    tree_view = QTreeView()
    tree_view.setModel(model)
    vlayout.addWidget(tree_view)

    hlayout = QHBoxLayout()

    data = QPushButton()
    data.setText('Change Data')
    data.clicked.connect(model.data_clicked)
    hlayout.addWidget(data)

    single = QPushButton()
    single.setText('Update Single')
    single.clicked.connect(model.single_clicked)
    hlayout.addWidget(single)

    double = QPushButton()
    double.setText('Update Double')
    double.clicked.connect(model.double_clicked)
    hlayout.addWidget(double)

    vlayout.addLayout(hlayout)

    window.show()
    return app.exec_()


if __name__ == '__main__':
    sys.exit(main())