Get currently selected cell of QTreeWidget

I'd like to modify QTreeWidget to make the selected cell editable when the enter key is hit, but keep the selection to full rows.

I've done a hacky implementation of figuring out where the last click was and saving the value, then sending those values to my edit_item function on the key press (also used for the itemDoubleClicked signal). It's not great though and I'm wondering if there's a much easier way to do it.

For the record, clicking on an item still selects the whole row. It's probably hidden behaviour by default, but in Maya there's a visible selection thing of the last cell that was moved over while the mouse button was held. If I could somehow get access to that, I could also add in behaviour to control it with the arrow keys.

This is an example of the selected cell:

This is my code so far:

class QTreeWidget(QtWidgets.QTreeWidget):

    returnPressed = QtCore.Signal(QTreeWidget, int)

    def __init__(self,  *args, **kwargs):
        QtWidgets.QTreeWidget.__init__(self, *args, **kwargs)

    def keyPressEvent(self, event):
        if event.key() == QtCore.Qt.Key_Return:
            self.returnPressed.emit(self._selected_item, self._selected_column)
        else:
            QtWidgets.QTreeWidget.keyPressEvent(self, event)

    def _mouse_pos_calculate(self, x_pos):
        """Find the currently selected column."""
        try:
            item = self.selectedItems()[0]
        except IndexError:
            item = None

        header = self.header()
        total_width = 0
        for i in range(self.columnCount()):
            total_width += header.sectionSize(i)
            if total_width > x_pos:
                return (item, i)

    def mousePressEvent(self, event):
        QtWidgets.QTreeWidget.mousePressEvent(self, event)
        self._selected_item, self._selected_column = self._mouse_pos_calculate(event.pos().x())

    def mouseReleaseEvent(self, event):
        QtWidgets.QTreeWidget.mouseReleaseEvent(self, event)
        self._selected_item, self._selected_column = self._mouse_pos_calculate(event.pos().x())

Edit: Improved function thanks to eyllanesc

class QTreeWidget(QtWidgets.QTreeWidget):
    """Add ability to edit cells when pressing return."""
    itemEdit = QtCore.Signal(QtWidgets.QTreeWidgetItem, int)

    def __init__(self,  *args, **kwargs):
        QtWidgets.QTreeWidget.__init__(self, *args, **kwargs)

        self._last_item = None
        self._last_column = 0
        self.itemDoubleClicked.connect(self._edit_item_intercept)

    def _edit_item_intercept(self, item=None, column=None):
        if item is None:
            item = self._last_item
        if column is None:
            column = self._last_column
        self.itemEdit.emit(item, column)

    def _store_last_cell(self, pos):
        selected_item = self.itemAt(pos)
        if selected_item is None:
            return
        self._last_item = selected_item
        self._last_column = self.header().logicalIndexAt(pos.x())

    def keyPressEvent(self, event):
        if event.key() == QtCore.Qt.Key_Return:
            return self._edit_item_intercept()
        QtWidgets.QTreeWidget.keyPressEvent(self, event)

    def mouseMoveEvent(self, event):
        QtWidgets.QTreeWidget.mouseMoveEvent(self, event)
        self._store_last_cell(event.pos())

Solution

You are doing a lot of calculation unnecessarily, in the next part I show a cleaner solution:

from PySide2 import QtCore, QtGui, QtWidgets


class QTreeWidget(QtWidgets.QTreeWidget):
    def __init__(self, *args, **kwargs):
        super(TreeWidget, self).__init__(*args, **kwargs)
        self.special_item = None
        self.special_col = 0

    def keyPressEvent(self, event):
        if event.key() == QtCore.Qt.Key_Return:
            self.editItem(self.special_item, self.special_col)
        QtWidgets.QTreeWidget.keyPressEvent(self, event)

    def editEnable(self, pos):
        press_item = self.itemAt(pos)
        if press_item is None:
            return
        if press_item is self.selectedItems()[0]:
            col = self.header().logicalIndexAt(pos.x())
            self.special_item = press_item
            self.special_col = col

    def mousePressEvent(self, event):
        QtWidgets.QTreeWidget.mousePressEvent(self, event)
        self.editEnable(event.pos())