drag from qtableview and drop on maya shelf to insert shelf button

test_drop_shelf_button.py

from __future__ import division

import maya.OpenMayaUI
import maya.cmds
from pymel.core import melGlobals


from PySide2 import QtCore, QtGui, QtWidgets
import shiboken2

import six

if six.PY3:
    long = int


def get_qobj(ui, widget=QtCore.QObject):
    if isinstance(ui, six.string_types):
        ptr = maya.OpenMayaUI.MQtUtil.findControl(ui)
        if ptr is None:
            ptr = maya.OpenMayaUI.MQtUtil.findLayout(ui)
        if ptr is None:
            ptr = maya.OpenMayaUI.MQtUtil.findMenuItem(ui)
        if ptr is not None:
            return shiboken2.wrapInstance(long(ptr), widget)
    else:
        return shiboken2.wrapInstance(long(ui), widget)


def get_maya_ui(qobj):
    return maya.OpenMayaUI.MQtUtil.fullName(long(shiboken2.getCppPointer(qobj)[0]))


class TableModel(QtCore.QAbstractTableModel):
    def __init__(self, data):
        super(TableModel, self).__init__()
        self._data = data

    def data(self, index, role):
        if role == QtCore.Qt.DisplayRole:
            return self._data[index.row()][index.column()]

    def rowCount(self, index):
        return len(self._data)

    def columnCount(self, index):
        return len(self._data[0])

    def flags(self, index):
        return (
            QtCore.Qt.ItemIsEnabled
            | QtCore.Qt.ItemIsSelectable
            | QtCore.Qt.ItemIsDragEnabled
        )


class Table(QtWidgets.QTableView):
    dropped = QtCore.Signal(QtCore.QPoint, list)

    def __init__(self, parent=None):
        super(Table, self).__init__(parent)
        self.setSelectionBehavior(QtWidgets.QAbstractItemView.SelectRows)
        self.setDragDropMode(QtWidgets.QAbstractItemView.DragDrop)
        self.setDropIndicatorShown(True)

    def startDrag(self, supportedActions):
        listsQModelIndex = self.selectedIndexes()
        if listsQModelIndex:
            dragQDrag = QtGui.QDrag(self)
            defaultDropAction = QtCore.Qt.CopyAction
            dataQMimeData = self.model().mimeData(listsQModelIndex)
            dragQDrag.setMimeData(dataQMimeData)
            # dragQDrag.targetChanged.connect(self.__test)
            print("drag start")
            dragQDrag.exec_(supportedActions, defaultDropAction)
            pos = QtGui.QCursor.pos()
            data = [x.data() for x in listsQModelIndex]
            n = 2
            data = [data[i : i + n] for i in range(0, len(data), n)]
            self.dropped.emit(pos, data)


class MainWindow(QtWidgets.QMainWindow):
    def __init__(self, parent=None):
        super(MainWindow, self).__init__(parent)

        self.table = Table()
        self.table.horizontalHeader().setStretchLastSection(True)

        data = [
            ["awesome", "1.0"],
            ["ordinary", "0.3"],
            ["good", "1.2"],
            ["bad", "0.1"],
        ]

        self.model = TableModel(data)
        self.table.setModel(self.model)
        self.setCentralWidget(self.table)

        self.table.dropped.connect(self.insertButton)

    def insertButton(self, pos, data):
        print('insert at %s' % pos, data) 
        shelfTabLayout = melGlobals["gShelfTopLevel"]
        shelfTabLayoutO = get_qobj(shelfTabLayout, QtWidgets.QWidget)
        pos_local = shelfTabLayoutO.mapFromGlobal(pos)

        maya.cmds.shelfTabLayout(shelfTabLayout, query=True, selectTabIndex=True)
        widget = QtWidgets.QApplication.widgetAt(pos)
        ui = get_maya_ui(widget)
        ui = ui.rstrip("|")
        print('dropped on', ui)
        current = maya.cmds.shelfTabLayout(shelfTabLayout, q=1, st=1)
        if ui.endswith("%s|mayaLayoutInternalWidget" % current):
            move = False
        else:
            ui_type = maya.cmds.objectTypeUI(ui)
            if ui_type in ("shelfButton", "separator"):
                move = True
            elif ui_type == "tabLayout":
                move = False
                # figure out which tab it is
                tabWidget = widget.parent()
                bar = tabWidget.tabBar()
                shelves = len(
                    maya.cmds.shelfTabLayout(shelfTabLayout, query=True, ca=True)
                )
                tab_width = bar.width() / shelves

                # estimated tab
                index = int((pos_local.x() // tab_width)) + 1
                print(
                    "dropped on",
                    index,
                    maya.cmds.shelfTabLayout(shelfTabLayout, q=1, tabLabelIndex=1)[
                        index
                    ],
                )
                # now you can append shelf button to that shelf tab and
                # activate it
                return
            else:
                return

        # append shelf button to currente tab
        buttons = []
        for label, ver in data:
            button = maya.cmds.shelfButton(
                label=label,
                annotation="I am %s" % label,
                p="%s|%s" % (shelfTabLayout, current),
                command="print(%r)" % label,
                image1=":/ambientlight.png",
                imageOverlayLabel=label,
            )
            buttons.append(button)
        if move:
            # determine drop on left or right button under cursor drop pos
            width = maya.cmds.control(ui, q=1, w=1)
            # left =  pos_local.x() // width
            existing = maya.cmds.shelfLayout(current, q=1, ca=1)
            index = existing.index(ui.split("|")[-1])
            left = sum([maya.cmds.control(x, q=1, w=1) for x in existing[:index] if maya.cmds.control(x, q=1, vis=1)])
            right = 1 if (pos_local.x() - left) > width / 2.0 else 0
            index = index + 1 + right
            for b in reversed(buttons):
                maya.cmds.shelfLayout(current, e=1, position=[b, index])


window = MainWindow()
window.show()