from PyQt4 import QtCore, QtGui

class WidgetComboBoxTree(QtGui.QTreeView):

  def __init__(self, parent=None):
    QtGui.QTreeView.__init__(self, parent)
    self.setUniformRowHeights(True)
    self.setRootIsDecorated(True)
    self.setItemsExpandable(True)
    self.setSortingEnabled(True)
    self.setSelectionBehavior(QtGui.QTreeView.SelectRows)
    self.header().setVisible(False)
    self.setMinimumHeight(150)

class WidgetTreeCombo(QtGui.QComboBox):
  __curItem = None
  __tree = None

  def __init__(self, parent=None):
    QtGui.QComboBox.__init__(self)
    self.__tree = WidgetComboBoxTree()
    self.setView(self.__tree)

  def hidePopup(self):
    self.__setCurModInd(self.view().currentIndex())
    return QtGui.QComboBox.hidePopup(self)

  def showPopup(self):
    QtGui.QComboBox.showPopup(self)
    self.setRootModelIndex(QtCore.QModelIndex())
    self.__tree.expandAll()

  def currentModelIndex(self):
    return self.__curItem

  def setModel(self, model, sortOrder=QtCore.Qt.AscendingOrder):
    if sortOrder in (QtCore.Qt.AscendingOrder, QtCore.Qt.DescendingOrder):
      self.__tree.header().setSortIndicator(0, sortOrder)
    QtGui.QComboBox.setModel(self, model)

  def setCurrentModelIndex(self, modelIndex):
    self.__setCurModInd(modelIndex)
    self.setCurrentIndex(modelIndex.isValid() and modelIndex.row() or 0)

  def tree(self):
    return self.__tree

  def __setCurModInd(self, idx):
    if idx and idx.isValid:
      self.__curItem = QtCore.QPersistentModelIndex(idx)
      self.setRootModelIndex(self.__curItem.parent())
    else:
      self.__curItem = None
      self.setRootModelIndex(QtCore.QModelIndex())