source: sasview/src/sas/qtgui/MainWindow/CategoryManager.py @ fb560d2

import json
import os

from PyQt5 import QtGui
from PyQt5 import QtCore
from PyQt5 import QtWidgets

from collections import defaultdict
from sas.qtgui.Utilities.CategoryInstaller import CategoryInstaller
from sasmodels.sasview_model import load_standard_models

from .UI.CategoryManagerUI import Ui_CategoryManagerUI
from .UI.ChangeCategoryUI import Ui_ChangeCategoryUI

class ToolTippedItemModel(QtGui.QStandardItemModel):
    """
    Subclass from QStandardItemModel to allow displaying tooltips in
    QTableView model.
    """
    def __init__(self, parent=None):
        QtGui.QStandardItemModel.__init__(self, parent)

    def headerData(self, section, orientation, role=QtCore.Qt.DisplayRole):
        """
        Displays tooltip for each column's header
        :param section:
        :param orientation:
        :param role:
        :return:
        """
        if role == QtCore.Qt.ToolTipRole:
            if orientation == QtCore.Qt.Horizontal:
                return str(self.header_tooltips[section])

        return QtGui.QStandardItemModel.headerData(self, section, orientation, role)

class Categories(object):
    """
    Container class for accessing model categories
    """
    def __init__(self):
        self.master_category_dict = defaultdict(list)
        self.by_model_dict = defaultdict(list)
        self.model_enabled_dict = defaultdict(bool)
        self.models = {}

        # Prepare the master category list
        self.readCategoryInfo()

        # Prepare model->category lookup
        self.setupModelDict()

    def readCategoryInfo(self):
        """
        Reads the categories in from file
        """
        categorization_file = CategoryInstaller.get_user_file()
        if not os.path.isfile(categorization_file):
            categorization_file = CategoryInstaller.get_default_file()
        with open(categorization_file, 'rb') as cat_file:
            self.master_category_dict = json.load(cat_file)
            self.regenerateModelDict()

        self.category_list = sorted(self.master_category_dict.keys())

    def saveCategories(self):
        """
        Serializes categorization info to file
        """
        self.regenerateMasterDict()
        with open(CategoryInstaller.get_user_file(), 'w') as cat_file:
            json.dump(self.master_category_dict, cat_file)

    def setupModelDict(self):
        """
        create a dictionary for model->category mapping
        """
        # Load the model dict
        models = load_standard_models()
        for model in models:
            # {model name -> model object}
            self.models[model.name] = model

    def regenerateModelDict(self):
        """
        Regenerates self.by_model_dict which has each model name as the
        key and the list of categories belonging to that model
        along with the enabled mapping
        """
        self.by_model_dict = defaultdict(list)
        for category in self.master_category_dict:
            for (model, enabled) in self.master_category_dict[category]:
                self.by_model_dict[model].append(category)
                self.model_enabled_dict[model] = enabled

    def regenerateMasterDict(self):
        """
        regenerates self.master_category_dict from
        self.by_model_dict and self.model_enabled_dict
        """
        self.master_category_dict = defaultdict(list)
        for model in self.by_model_dict:
            for category in self.by_model_dict[model]:
                self.master_category_dict[category].append\
                    ((model, self.model_enabled_dict[model]))

    def modelToCategory(self):
        """
        Getter for the model->category dict
        """
        return self.by_model_dict

    def modelDict(self):
        """
        Getter for the model list
        """
        return self.models

    def categoryDict(self):
        """
        Getter for the category dict
        """
        return self.master_category_dict

    def categoryList(self):
        """
        Getter for the category list
        """
        return self.category_list



class CategoryManager(QtWidgets.QDialog, Ui_CategoryManagerUI):
    def __init__(self, parent=None, manager=None):
        super(CategoryManager, self).__init__(parent)
        self.setupUi(self)

        self.communicator = manager.communicator()

        self.setWindowTitle("Category Manager")

        self.initializeGlobals()

        self.initializeModels()

        self.initializeSignals()

    def initializeGlobals(self):
        """
        Initialize global variables used in this class
        """
        # Default checked state
        self.chkEnable.setCheckState(QtCore.Qt.Checked)

        # Modify is disabled by default
        self.cmdModify.setEnabled(False)

        # Data for chosen model
        self.model_data = None

        # Categories object
        self.categories = Categories()

    def initializeModels(self):
        """
        Set up models and views
        """
        # Set the main models
        self._category_model = ToolTippedItemModel()
        self.lstCategory.setModel(self._category_model)
        # Proxy model for showing a subset of model content
        self.model_proxy = QtCore.QSortFilterProxyModel(self)
        self.model_proxy.setSourceModel(self._category_model)
        self.lstCategory.setModel(self.model_proxy)

        self.initializeModelList()

        self.setTableProperties(self.lstCategory)

        self.lstCategory.setAlternatingRowColors(True)
        # self.lstCategory.setContextMenuPolicy(QtCore.Qt.CustomContextMenu)
        # self.lstCategory.setAttribute(QtCore.Qt.WA_MacShowFocusRect, False)

    def initializeModelList(self):
        """
        Model category combo setup
        """
        self._category_model.clear()
        for ind, model in enumerate(self.categories.modelDict()):
            item = QtGui.QStandardItem(model)
            empty_item = QtGui.QStandardItem()
            empty_item.setEditable(False)
            # Add a checkbox to it
            item.setCheckable(True)
            item.setCheckState(QtCore.Qt.Checked)
            item.setEditable(False)
            current_category = self.categories.modelToCategory()[model]
            self._category_model.appendRow([item, empty_item])
            self._category_model.item(ind, 1).setText(', '.join(i for i in current_category))

    def initializeSignals(self):
        """
        :return:
        """
        self.cmdOK.clicked.connect(self.onClose)
        self.cmdModify.clicked.connect(self.onModify)
        self.cmdReset.clicked.connect(self.onReset)

        self.chkEnable.toggled.connect(self.onEnableAll)

        # every change in txtSearch
        self.txtSearch.textChanged.connect(self.onSearch)

        # Signals from the list
        selectionModel = self.lstCategory.selectionModel()
        selectionModel.selectionChanged.connect(self.onListSelection)


    def onClose(self):
        """
        Save the category file before exiting
        """
        self.categories.saveCategories()
        # Ask the fitting widget to reload the comboboxes
        self.communicator.updateModelCategoriesSignal.emit()

        self.close()

    def selectedModels(self):
        """
        Returns a list of selected models
        """
        selected_models = []
        selectionModel = self.lstCategory.selectionModel()
        selectedRows = selectionModel.selectedRows()
        for row in selectedRows:
            model_index = self.model_proxy.mapToSource(row)
            current_text = self._category_model.itemFromIndex(model_index).text()
            selected_models.append(self.categories.modelDict()[current_text])
        return selected_models

    def onListSelection(self):
        """
        Respond to row selection and update GUI
        """
        selected_items = self.selectedModels()
        self.cmdModify.setEnabled(len(selected_items) == 1)

    def onReset(self):
        """
        Reload the saved categories
        """
        self.initializeGlobals()
        # Reload the Categories object
        self.categories = Categories()
        # Reload the model
        self.initializeModelList()
        self.setTableProperties(self.lstCategory)
        self.lstCategory.setAlternatingRowColors(True)

    def onEnableAll(self, isChecked):
        """
        Respond to the Enable/Disable All checkbox
        """
        select = QtCore.Qt.Checked if isChecked else QtCore.Qt.Unchecked
        for row in range(self._category_model.rowCount()):
            self._category_model.item(row).setCheckState(select)

    def onSearch(self):
        """
        Respond to text entered in search field
        """
        if self.txtSearch.isModified():
            input_to_check = str(self.txtSearch.text())

        # redefine the proxy model
        self.model_proxy.setFilterRegExp(QtCore.QRegExp(input_to_check,
                         QtCore.Qt.CaseInsensitive, QtCore.QRegExp.FixedString))

    def onModify(self):
        """
        Show the Change Category dialog - modal
        """
        # Selected category:
        selected_models = self.selectedModels()
        if len(selected_models) > 1:
            # Somehow we got more than one model - complain!
            raise AttributeError("Please select only one model.")
        change_dialog = ChangeCategory(parent=self, categories=self.categories, model=selected_models[0])

        if change_dialog.exec_() != QtWidgets.QDialog.Accepted:
            return
        # Reload the model
        self.initializeModelList()
        self.setTableProperties(self.lstCategory)
        self.lstCategory.setAlternatingRowColors(True)

    def setTableProperties(self, table):
        """
        Setting table properties
        """
        # Table properties
        table.verticalHeader().setVisible(False)
        table.setAlternatingRowColors(True)
        #table.setSizePolicy(QtWidgets.QSizePolicy.MinimumExpanding, QtWidgets.QSizePolicy.Expanding)
        table.setSizePolicy(QtWidgets.QSizePolicy.Expanding, QtWidgets.QSizePolicy.Expanding)
        table.setSelectionBehavior(QtWidgets.QAbstractItemView.SelectRows)
        table.resizeColumnsToContents()

        # Header
        header = table.horizontalHeader()
        header.setSectionResizeMode(QtWidgets.QHeaderView.ResizeToContents)
        header.ResizeMode(QtWidgets.QHeaderView.Interactive)


        self._category_model.setHeaderData(0, QtCore.Qt.Horizontal, "Model")
        self._category_model.setHeaderData(1, QtCore.Qt.Horizontal, "Category")

        self._category_model.header_tooltips = ['Select model',
                                                'Change or create new category']

class ChangeCategory(QtWidgets.QDialog, Ui_ChangeCategoryUI):
    """
    Dialog for adding/removing categories for a single model
    """
    def __init__(self, parent=None, categories=None, model=None):
        super(ChangeCategory, self).__init__(parent)
        self.setupUi(self)

        self.model = model
        self.parent = parent
        self.categories = categories

        self.initializeElements()

        self.initializeList()

        self.initializeSignals()

    def initializeElements(self):
        """
        Initialize local GUI elements with information from the Categories object
        """
        self.cbCategories.addItems(self.categories.categoryList())
        self.setWindowTitle("Change Category for: "+ self.model.name)
        self.lblTitle.setText("Current categories for " + self.model.name)
        self.rbExisting.setChecked(True)
        self.onAddChoice()

    def initializeList(self):
        """
        Initialize the category list for the given model
        """
        current_categories = self.categories.modelToCategory()[self.model.name]
        for cat in current_categories:
            self.lstCategories.addItem(cat)

    def initializeSignals(self):
        """
        Initialize signals for UI elements
        """
        self.cmdAdd.clicked.connect(self.onAdd)
        self.cmdRemove.clicked.connect(self.onRemove)
        self.cmdOK.clicked.connect(self.onOK)
        self.cmdCancel.clicked.connect(self.close)

        # Signals from the list
        self.lstCategories.itemSelectionChanged.connect(self.onListSelection)

        # Signals from the radio buttons
        self.rbExisting.toggled.connect(self.onAddChoice)

    def onAddChoice(self):
        """
        Respond to the type selection for new category
        """
        isNew = self.rbNew.isChecked()
        self.cbCategories.setEnabled(not isNew)
        self.txtNewCategory.setEnabled(isNew)

    def onListSelection(self):
        """
        Respond to selection in the category list view
        """
        selected_items = self.selectedModels()
        self.cmdRemove.setEnabled(len(selected_items) > 0)

    def selectedModels(self):
        """
        Returns a list of selected models
        """
        selected_categories = []
        selectedRows = self.lstCategories.selectedItems()
        selected_categories = [str(row.text()) for row in selectedRows]

        return selected_categories

    def onAdd(self):
        """
        Add the chosen category to the list
        """
        if self.rbExisting.isChecked():
            new_category = self.cbCategories.currentText()
        else:
            new_category = self.txtNewCategory.text()
        # Display the current value as txt
        if new_category:
            self.lstCategories.addItem(new_category)

    def onRemove(self):
        """
        Remove selected categories in the list
        """
        selectedRows = self.lstCategories.selectedItems()
        if not selectedRows:
            return
        for row in selectedRows:
            self.lstCategories.takeItem(self.lstCategories.row(row))

    def onOK(self):
        """
        Accept the new categories for the model
        """
        # Read in the categories
        self.categories.modelToCategory()[self.model.name] = self.listCategories()
        self.accept()

    def listCategories(self):
        """
        Returns the list of categories from the QListWidget
        """
        return [str(self.lstCategories.item(i).text()) for i in range(self.lstCategories.count())]