qutebrowser/qutebrowser/utils/completion.py

from collections import OrderedDict

from PyQt5.QtCore import (QAbstractItemModel, Qt, QModelIndex, QVariant,
                          QSortFilterProxyModel, pyqtSignal)

class CompletionModel(QAbstractItemModel):
    def __init__(self, parent=None):
        super().__init__(parent)
        self._data = OrderedDict()
        self.parents = []
        self.root = CompletionItem([""] * 2)

    def removeRows(self, position=0, count=1, parent=QModelIndex()):
        node = self.node(parent)
        self.beginRemoveRows(parent, position, position + count - 1)
        node.children.pop(position)
        self.endRemoveRows()

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

    def columnCount(self, parent=QModelIndex()):
        # pylint: disable=unused-argument
        return self.root.column_count()

    def data(self, index, role=Qt.DisplayRole):
        if not index.isValid():
            return QVariant()
        item = index.internalPointer()
        try:
            return QVariant(item.data(index.column(), role))
        except (IndexError, ValueError):
            return QVariant()

    def flags(self, index):
        # FIXME categories are not selectable, but moving via arrow keys still
        # tries to select them
        if not index.isValid():
            return Qt.NoItemFlags
        flags = Qt.ItemIsEnabled
        if len(index.internalPointer().children) > 0:
            return flags
        else:
            return flags | Qt.ItemIsSelectable

    def headerData(self, section, orientation, role=Qt.DisplayRole):
        if orientation == Qt.Horizontal and role == Qt.DisplayRole:
            return QVariant(self.root.data(section))
        return None

    def setData(self, index, value, role=Qt.EditRole):
        if not index.isValid():
            return False
        item = index.internalPointer()
        try:
            item.setdata(index.column(), value, role)
        except (IndexError, ValueError):
            return False
        self.dataChanged.emit(index, index)
        return True

    def index(self, row, column, parent=QModelIndex()):
        if (0 <= row < self.rowCount(parent) and
                0 <= column < self.columnCount(parent)):
            pass
        else:
            return QModelIndex()

        if not parent.isValid():
            parent_item = self.root
        else:
            parent_item = parent.internalPointer()

        child_item = parent_item.children[row]
        if child_item:
            return self.createIndex(row, column, child_item)
        else:
            return QModelIndex()

    def parent(self, index):
        if not index.isValid():
            return QModelIndex()
        item = index.internalPointer().parent
        if item == self.root or item is None:
            return QModelIndex()
        return self.createIndex(item.row(), 0, item)

    def rowCount(self, parent=QModelIndex()):
        if parent.column() > 0:
            return 0

        if not parent.isValid():
            pitem = self.root
        else:
            pitem = parent.internalPointer()

        return len(pitem.children)

    def sort(self, column, order=Qt.AscendingOrder):
        raise NotImplementedError

    def init_data(self):
        for (cat, items) in self._data.items():
            newcat = CompletionItem([cat], self.root)
            self.root.children.append(newcat)
            for item in items:
                newitem = CompletionItem(item, newcat)
                newcat.children.append(newitem)

class CompletionItem():
    parent = None
    _data = None
    children = None
    _marks = None

    def __init__(self, data, parent=None):
        self.parent = parent
        self._data = data
        self.children = []
        self._marks = []

    def data(self, column, role=Qt.DisplayRole):
        if role == Qt.DisplayRole:
            return self._data[column]
        elif role == Qt.UserRole:
            return self._marks
        else:
            raise ValueError

    def setdata(self, column, value, role=Qt.DisplayRole):
        if role == Qt.DisplayRole:
            self._data[column] = value
        elif role == Qt.UserRole:
            self._marks = value
        else:
            raise ValueError

    def column_count(self):
        return len(self._data)

    def row(self):
        if self.parent:
            return self.parent.children.index(self)
        return 0

class CompletionFilterModel(QSortFilterProxyModel):
    _pattern = None
    pattern_changed = pyqtSignal(str)

    @property
    def pattern(self):
        return self._pattern

    @pattern.setter
    def pattern(self, val):
        self._pattern = val
        self.invalidate()
        self.pattern_changed.emit(val)

    def __init__(self, parent=None):
        super().__init__(parent)
        self.pattern = ''

    def filterAcceptsRow(self, row, parent):
        if parent == QModelIndex():
            return True
        idx = self.sourceModel().index(row, 0, parent)
        data = self.sourceModel().data(idx).value()
        # TODO more sophisticated filtering
        if not self.pattern:
            return True
        return self.pattern in data

    def lessThan(self, lindex, rindex):
        left = self.sourceModel().data(lindex).value()
        right = self.sourceModel().data(rindex).value()

        leftstart = left.startswith(self.pattern)
        rightstart = right.startswith(self.pattern)

        if leftstart and rightstart:
            return left < right
        elif leftstart:
            return True
        elif rightstart:
            return False
        else:
            return left < right