PyQt4 QAbstractItemModel subclasses

item_data.py

""" Define a data container to be used with model subclasses that emits a generic signal whenever data in the internal dictionary 
    is changed. This creates a consistent API layer with model items that can be edited programatically as dictionaries, and
    automatically kept synchronized in the model and the view.
"""
__author__ = 'Ian Davis'
__all__ = ['ItemData', ]

from api.util.event_util import Signal


class ItemData(object):
    """ Generic PyQt Model data container, that uses an internal signal class to avoid QObject limitations
        but still emit a signal when data is changed on the object.
    """

    def __init__(self, data):
        """ ItemData initializer.

        :param data: The dictionary of data.
        """
        self._data = data
        self.changed = Signal()

    def __contains__(self, key):
        return key in self._data

    def __getitem__(self, item):
        return self._data[item]

    def __setitem__(self, key, value):
        self._data[key] = value
        self.changed.emit()

    def iteritems(self):
        return self._data.iteritems()

    def iterkeys(self):
        return self._data.iterkeys()

    def itervalues(self):
        return self._data.itervalues()

    def __str__(self):
        return str(self._data)

table_model.py

""" Define a QTableModel subclass that uses dictionaries instead of column indexes and maps them to an internal header list to manage data.
"""
__author__ = 'Ian Davis'
__all__ = ['TableRow', 'TableModel', ]

import re
from collections import OrderedDict as OrderedDictionary

from PyQt4.QtCore import QAbstractTableModel
from PyQt4.QtCore import QModelIndex
from PyQt4.QtCore import QObject
from PyQt4.QtCore import Qt

from api.models import ItemData
from api.util.event_util import Signal


class TableRow(object):
    """ TableModel data container, represents one row in the table rather than a single item.

        Useful for most cases as a table's rows are rigidly structured, and columns are the same.
    """
    def __init__(self, data, row=0):
        self.row = row
        self.data = ItemData(data)
        self.changed = Signal()

        self._connect_slots()

    def _connect_slots(self):
        self.data.changed.connect(self.changed.emit)

    def __getitem__(self, item):
        return self.data[item]

    def __setitem__(self, item, value):
        self.data[item] = value

    def iteritems(self):
        return self.data.iteritems()

    def iterkeys(self):
        return self.data.iterkeys()

    def itervalues(self):
        return self.data.itervalues()


class TableModel(QAbstractTableModel):
    """ TableModel is an implementation of PyQt's QAbstractTableModel that overrides default indexing to use dictionary key-based mapping,
        mapping a column in the table's header to a value for that column. The goal here is to simplify indexing by being able to manage
        the data in a table based on string keys instead of arbitrary indexes, eliminating the need to cross-reference a header to find where
        to put a value.
    """
    def __init__(self, header, header_types=None, key_column=None, parent=None):
        """ TableModel initializer

        :param header: A list containing the header values for the table.
        :param header_types: A dictionary mapping the header values to their types, default is string.
        :param key_column: The primary key column for the table (the column to reference rows by).z
        :param parent: The QT Parent widget.
        """
        QAbstractTableModel.__init__(self, parent)

        self.header = header
        self.header_types = header_types

        if not self.header_types:
            self.header_types = {}
            for column in self.header:
                self.header_types[column] = 'string'

        self.key_column = key_column

        if not self.key_column:
            self.key_column = self.header[0]

        self.data = OrderedDictionary()

    def rowCount(self, parent=QModelIndex()):
        """ Model-method, called by the view to determine how many rows are to be displayed at a given time.
        """
        return len(self.data)

    def columnCount(self, parent=QModelIndex()):
        """ Model-method, called by the view to determine how many columns are to be displayed at a given time.
        """
        return len(self.header)

    def setHeaderData(self, section, orientation, role):
        """ Called to set the data for a given column in the header.

        :param section: The header section to change.
        :param orientation: The orientation of the section (Horizontal or Vertical).
        :param role: The role of the section (DisplayRole, etc).
        """
        self.headerDataChanged.emit(orientation, section, section)
        return True

    def headerData(self, section, orientation, role):
        """ Model-method, called by the view to determine what to display for a given section of the header.

        :param section: The section to display
        :param orientation: The orientation of the section (Horizontal or Vertical).
        :param role: The role of the section (DisplayRole, etc).
        :return:
        """
        if orientation == Qt.Horizontal and role == Qt.DisplayRole:
            if section >= len(self.header):
                return

            return self.header[section]

    def index(self, row, col, parent=QModelIndex()):
        """ Model-method, Return a QModelIndex that points to a given row, column and parent (parent is for Tree-based models mainly).

            Uses internal method createIndex defined by Qt to create a QModelIndex instance.

        :param row: The row of this index.
        :param col: The column of this index.
        :param parent: The parent of this index.
        :return: QModelIndex pointing at the given row and column.
        """
        table_row = self.data.values()[row]
        return self.createIndex(row, col, table_row)

    def setData(self, index, data, role):
        """ Model-method, called by the view when a given index's data is changed to update the model with that change.

            In here, we lookup the pointer from the index (which will be an instance of our internal TableRow class),
            get the column name for the column edited, and set the table row's dictionary value for that column to the data entered.

        :param index:
        :param data:
        :param role:
        :return:
        """
        if not index.isValid():
            return False

        elif index.column() >= len(self.header):
            return

        elif not role == Qt.EditRole:
            return False

        table_row = index.internalPointer()
        column_name = self.header[index.column()]
        table_row[column_name] = str(data.toString())

        self.dataChanged.emit(index, index)

        return True

    def data(self, index, role):
        """ Model-method, called by the view to determine what to display for a given index and role.

        :param index: QModelIndex to display data for.
        :param role: The role to display (DisplayRole, TextAlignmentRole, etc).
        :return: The data to display.
        """
        if not index.isValid():
            return

        elif index.column() >= len(self.header):
            return

        elif role == Qt.TextAlignmentRole:
            return Qt.AlignCenter

        elif not role == Qt.DisplayRole:
            return

        table_row = index.internalPointer()
        column_name = self.header[index.column()]
        data = table_row[column_name]

        if not isinstance(data, QObject):
            if not data:
                data = ''

            data = str(data)

        return data

    def flags(self, index):
        """ QAbstractTableModel override method that is used to set the flags for the item at the given QModelIndex.
            Here, we just set all indexes to enabled, and selectable.
        """
        return Qt.ItemIsEnabled | Qt.ItemIsSelectable | Qt.ItemIsEditable

    def add_row(self, data):
        """ Add a new row to the table, displaying the data mapped from a dictionary to our table header.

        :param data: A dictionary mapping to our table header and the values for each column.
        :return: TableRow instance that was added to the model.
        """
        row = self.rowCount()
        table_row = TableRow(data, row)
        key_value = data[self.key_column]

        self.beginInsertRows(QModelIndex(), row, row)
        self.data[key_value] = table_row
        self._connect_node(table_row)
        self.endInsertRows()

        return table_row

    def removeRows(self, row, count, parent=QModelIndex()):
        """ Model-method to remove a number of rows, starting a row.

        :param row: The row to begin removing from.
        :param count: The number of rows to remove.
        :param parent: The parent index of the row to begin from.
        :return: True if the rows were successfully removed.
        """
        self.beginRemoveRows(parent, row, row + count)
        new_data = self.data.copy()

        for key in new_data.keys()[row:row + count]:
            del self.data[key]

        self.endRemoveRows()

        return True

    def _connect_node(self, node):
        node.changed.connect(lambda: self._notify_data_changed(node))

    def _notify_data_changed(self, node):
        row = node.row
        top_left = self.createIndex(row, 0, node)
        bottom_right = self.createIndex(row, len(self.header), node)
        self.dataChanged.emit(top_left, bottom_right)

    def find_index(self, pointer):
        """ Helper method to find a QModelIndex that points to a given pointer.

        :param pointer: The TableRow to find a QModelIndex for.
        :return: QModelIndex, or None.
        """
        for index in self.persistentIndexList():
            if index.column() != 0:
                continue

            if index.internalPointer() == pointer:
                return index

    def match_pattern(self, section, pattern):
        """ Match a given regex pattern to the rows in our table, and create a list of rows that matched.

        :param section: The column in the table to match against.
        :param pattern: The regex pattern to match.
        :return: The list of rows that matched the pattern.
        """
        compiled_regex = re.compile(pattern)
        column_name = self.header[section]
        rows_to_hide = []

        for table_row in self.data.itervalues():
            data = table_row[column_name]
            if not compiled_regex.match(str(data)):
                rows_to_hide.append(table_row.row)

        return rows_to_hide

    def pack_dictionary(self, dictionary):
        """ Given a dictionary, create a new dictionary with columns missing from the original replaced with empty strings.

        :param dictionary: The dictionary to pack.
        :return: The packed dictionary.
        """
        packed_dictionary = {}

        for column in self.header:
            packed_dictionary[column] = dictionary.get(column, '')

        return packed_dictionary

tree_model.py

""" Module Docstring
"""
__author__ = 'Ian Davis'
__all__ = ['TreeItem', 'TreeModel', ]

from collections import OrderedDict as OrderedDictionary

from PyQt4.QtCore import QAbstractItemModel
from PyQt4.QtCore import QModelIndex
from PyQt4.QtCore import QObject
from PyQt4.QtCore import Qt

from api.models import ItemData
from api.util.event_util import Signal


class TreeItem(object):
    """ TreeItem represents one node in a TreeModel instance, the main implementation difference from a regular
            QTreeModel is the use of a dictionary to map data to the models' columns instead of an ordered tuple/list.

        parent: The TreeItem instance that owns this instance.
        children: An OrderedDictionary, keyed by the key_column of the model, that contains the children TreeItems.
        data: A Dictionary used to represent the data that this TreeItem instance displays on the model.
    """

    def __init__(self, data, parent=None):
        """ TreeItem constructor.
        """
        self.data = ItemData(data)
        self.parent = parent
        self.children = OrderedDictionary()
        self.changed = Signal()

        self._connect_slots()
        self._initialized = True

    def _connect_slots(self):
        self.data.changed.connect(self.changed.emit)

    def row(self):
        """ This method is necessary because of the parent-child node structure of the model, where there is no simple
            way to find the overall relationship of all the items in the database, rather just one items' relationship
            with those surrounding it.
        :return: int
        """
        if not self.parent:
            return 0

        return self.parent.children.values().index(self)

    def __getitem__(self, item):
        return self.data[item]

    def __setitem__(self, item, value):
        self.data[item] = value

    def __iter__(self):
        return self.children.itervalues()

    def __str__(self):
        return '{0}({1}'.format(self.__class__.__name__, str(self.data))

    def __repr__(self):
        return str(self)


class TreeModel(QAbstractItemModel):
    """ TreeModel is an implementation of PyQt's QAbstractItemModel that overrides default indexing support to use
        python dictionaries mapping a column in the table header supplied to a value for said column. The goal here
        is to simplify indexing by being able to manage the data in a table based on string keys instead of arbitrary
        indexes, eliminating the need to cross-reference a header to find where to put a value.
    """

    def __init__(self, header, header_types=None, key_column=0, parent=None):
        """ TreeModel constructor
        :param header: The header to use
        :type header: Iterable
        :param parent: A QWidget that QT will give ownership of this Widget too.
        """
        super(TreeModel, self).__init__(parent)

        self.header = header
        self.header_types = header_types

        if not self.header_types:
            for column in self.header:
                self.header_types[column] = 'string'

        self.key_column = self.header[key_column]
        self.root = TreeItem(header)

    def find_index(self, pointer):
        for index in self.persistentIndexList():
            if index.column() != 0:
                continue

            if index.internalPointer() == pointer:
                return index

    def _connect_node(self, node):
        node.changed.connect(lambda: self._notify_data_changed(node))

    def _notify_data_changed(self, node):
        row = node.row()
        top_left = self.createIndex(row, 0, node)
        bottom_right = self.createIndex(row, len(self.header), node)
        self.dataChanged.emit(top_left, bottom_right)

    def add_node(self, values, children=None, parent=None):
        """ Add a new root TreeItem to our model, using the values passed as the data.
            Optional args: children, parent
        :param values: A dictionary mapping the model's header to the values to use for this TreeItem.
        :param children: A collection of dictionaries mapping the model's header to the values to use for each child
            TreeItem.
        :param parent: The parent to give ownership of this TreeItem too, if not given, defaults to the root TreeItem
        :return: The TreeItem instance that was added.
        """
        if not parent:
            parent = self.root

        key = values[self.key_column]
        node = TreeItem(values, parent)

        if children:
            for values_ in children:
                self.add_node(values_, parent=node)

        parent.children[key] = node

        self._connect_node(node)
        return node

    def remove_node(self, key_value, parent=None):
        """ Remove the node that matches the key value and parent given.
        :param key_value: str
        :param parent: TreeItem
        :return: bool
        """
        if not parent:
            parent = self.root

        parent_index = self.find_index(parent)
        if key_value not in parent.children:
            raise KeyError('{key} not found in {node}'.format(key=key_value, node=parent[self.key_column]))

        row = parent.children.keys().index(key_value)

        self.removeRow(row, parent_index)

        return True

    def flags(self, index):
        """ QAbstractItemModel override method that is used to set the flags for the item at the given QModelIndex.
            Here, we just set all indexes to enabled, and selectable.
        """
        return Qt.ItemIsEnabled | Qt.ItemIsSelectable

    def data(self, index, role):
        """ Return the data to display for the given index and the given role.
            This method should not be called directly. This method is called implicitly by the QTreeView that is
            displaying us, as the way of finding out what to display where.
        """
        if not index.isValid():
            return

        elif not role == Qt.DisplayRole:
            return

        item = index.internalPointer()
        column = self.header[index.column()]

        if column not in item.data:
            return

        data = item.data[column]

        if not isinstance(data, QObject):
            data = str(data)

        return data

    def index(self, row, col, parent):
        """ Return a QModelIndex instance pointing the row and column underneath the parent given.
            This method should not be called directly. This method is called implicitly by the QTreeView that is
            displaying us, as the way of finding out what to display where.
        """
        if not parent or not parent.isValid():
            parent = self.root

        else:
            parent = parent.internalPointer()

        if row < 0 or row >= len(parent.children.keys()):
            return QModelIndex()

        row_name = parent.children.keys()[row]

        child = parent.children[row_name]
        return self.createIndex(row, col, child)

    def parent(self, index=None):
        """ Return the index of the parent TreeItem of a given index. If index is not supplied, return an invalid
                QModelIndex.
            Optional args: index

        :param index: QModelIndex
        :return:
        """
        if not index:
            return QModelIndex()

        elif not index.isValid():
            return QModelIndex()

        child = index.internalPointer()
        parent = child.parent

        if parent == self.root:
            return QModelIndex()

        elif child == self.root:
            return QModelIndex()

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

    def rowCount(self, parent):
        """ Return the number of rows a given index has under it. If an invalid QModelIndex is supplied, return the
                number of children under the root.

        :param parent: QModelIndex
        """
        if parent.column() > 0:
            return 0

        if not parent.isValid():
            parent = self.root

        else:
            parent = parent.internalPointer()

        return len(parent.children)

    def columnCount(self, parent):
        """ Return the number of columns in the model header. The parent parameter exists only to support the signature
                of QAbstractItemModel.
        """
        return len(self.header)

    def headerData(self, section, orientation, role):
        """ Return the header data for the given section, orientation and role. This method should not be called
            directly. This method is called implicitly by the QTreeView that is displaying us, as the way of finding
            out what to display where.
        """
        if orientation == Qt.Horizontal and role == Qt.DisplayRole:
            return self.root.data[section]

    def __iter__(self):
        for child in self.root.children.itervalues():
            yield child