Source code for schrodinger.ui.qt.mapperwidgets._comboboxes

from schrodinger.models import mappers
from schrodinger.ui.qt import utils as qtutils
from schrodinger.Qt import QtCore
from schrodinger.Qt import QtGui
from schrodinger.Qt import QtWidgets

_INVALID_ITEM_TEXT = "ERROR: Model value not found in combobox"


class AbstractMappableComboBox(QtWidgets.QComboBox):
    """
    A combobox based on userData (a.k.a. item)
    """
    currentItemChanged = QtCore.pyqtSignal(object)

    def __init__(self, parent=None):
        super().__init__(parent)
        self.currentIndexChanged.connect(self._onCurrentIndexChanged)
        self._invalid_item_idx = None

    def _onCurrentIndexChanged(self, idx):
        self.currentItemChanged.emit(self.currentItem())

    def isValid(self):
        return self._invalid_item_idx is None

    def addItem(self, text, value=None, **kwargs):
        super().addItem(text, userData=value, **kwargs)
        if (self._invalid_item_idx is not None and
                self.itemData(self._invalid_item_idx) == value):
            self._clearInvalidItem()
            self.setCurrentItem(value)

    def addItems(self, items):
        """
        :param items: Dictionary mapping the desired text for the combo box
            items mapped to the values to set the model to.
        :type  items: dict[str, object]
        """
        for item in items.items():
            self.addItem(*item)

    def currentItem(self):
        return self.itemData(self.currentIndex())

    def clear(self):
        current_item = self.currentItem()
        with qtutils.suppress_signals(self):
            super().clear()
        self._invalid_item_idx = None
        self.setCurrentItem(current_item)

    def _clearInvalidItem(self):
        if self._invalid_item_idx is not None:
            self.removeItem(self._invalid_item_idx)
            self._invalid_item_idx = None

    def setCurrentIndex(self, index):
        if self.itemText(index) != _INVALID_ITEM_TEXT:
            if self._invalid_item_idx is not None and index > self._invalid_item_idx:
                index -= 1
                # Changing to a non invalid item so we can clear it
            self._clearInvalidItem()
        super().setCurrentIndex(index)

    def setCurrentItem(self, new_current_item):
        for idx in range(self.count()):
            item = self.itemData(idx)
            if new_current_item == item:
                self.setCurrentIndex(idx)
                return
        else:
            self._clearInvalidItem()
            invalid_item_idx = self.count()
            self.addItem(text=_INVALID_ITEM_TEXT, value=new_current_item)
            self.setItemData(invalid_item_idx, QtGui.QBrush(QtCore.Qt.red),
                             QtCore.Qt.ForegroundRole)
            self.setCurrentIndex(invalid_item_idx)
            self._invalid_item_idx = invalid_item_idx


class MappableComboBox(mappers.TargetMixin, AbstractMappableComboBox):
    """
    Simple extension of QComboBox that allows the combo box to be used with
    mappers.

    If this combobox is ever mapped to an item that hasn't been added to the
    combobox, then the combobox will have an item that reads  `ERROR: Model
    value not found in combobox`. If an item is added to the combobox
    that matches the model value or if the model value is changed to a valid
    value, then the ERROR item will be automatically removed.

    WARNING This class assumes that the data for each item is unique. If it is
    not, then the item will be chosen arbitrarily when the model is changed
    to the non-unique data value.
    """

    def __init__(self, parent=None):
        super().__init__(parent)
        self.currentItemChanged.connect(self.targetValueChanged)

    def targetGetValue(self):
        return self.currentItem()

    def targetSetValue(self, value):
        self.setCurrentItem(value)


class EnumComboBox(MappableComboBox):
    """
    A combo box meant to be used with an enum.
    """

    def __init__(self, parent=None, enum=None):
        super().__init__(parent=parent)
        self._listified_enum = None
        if enum is not None:
            self.setEnum(enum)

    def setEnum(self, enum):
        """
        Set the enum options as items of this combobox.

        :param enum: An enum to populate the items of this combobox
        :type enum: Enum or List(Enum)
        """
        # Save a list of the enum rather than just the enum for indexing
        # purposes.
        self._listified_enum = list(enum)
        self.clear()
        if all(isinstance(member.value, str) for member in list(enum)):
            for member in enum:
                self.addItem(member.value, member)
        else:
            for member in enum:
                self.addItem(member.name, member)
        self.setCurrentItem(self._listified_enum[0])

    def index(self, item):
        """
        :return: The index of the specified enum member
        """
        return self._listified_enum.index(item)

    def updateItemTexts(self, disp_text):
        """
        Update the item texts for this combo box. Expects a list of 2-tuples
        with the enum member to update the text and the text to update it to.

        For example, if the `EnumComboBox` was used with the enum
        `Enum('Foo','bar raf')`, this method could be called with
        `[(Foo.bar, 'FIRST_OPTION'), (Foo.raf, 'SECOND_OPTION')]`.
        """
        if isinstance(disp_text, list):
            for enum_member, text in disp_text:
                idx = self.index(enum_member)
                self.setItemText(idx, text)
        else:
            raise TypeError('Expected a list of tuples containing enum members '
                            'and what text should be displayed for them.')

    def setItemTexts(self, texts):
        """
        Update all item texts at once.

        :param texts: A list of strings, one per item.
        """
        if len(texts) != self.count():
            raise ValueError('Length of a list of strings must be '
                             ' equal to the number of enum members.')
        for idx, text in enumerate(texts):
            self.setItemText(idx, text)

    def setCurrentItem(self, new_current_item):
        """
        Sets the current item as the selected option.
        """
        super().setCurrentItem(new_current_item)
        if self._listified_enum is None:
            raise RuntimeError(
                'No enums have been set up with this EnumComboBox, '
                'use setEnum(enum) to setup the enum options.')