[Drag and drop between item widgets with custom types] normally drag-and-drop between item widgets doesn't preserve subclass types for custom items. This enables that, with interchangeable drag and drop between all three types of item widgets. #python #pyside #qt

rich_drag_and_drop.py

import abc
import json
import inspect
from collections import defaultdict
from typing import Type, Iterable
from functools import partial

from PySide2 import QtWidgets, QtGui, QtCore


class SerializableItemMixin:
    """
    This - alongside DragAndDropItemsConverter - allows for the drag and drop of
    copies of Qt ____WidgetItems between views while preserving their data types.

    If this is used multiple times, each should set its own SerializedTypeID value
    to ensure objects are reconstructed correctly. If using multiple ItemWidgets
    with the same serializable data types, they can use the same SerializedTypeID
    to enable drag and drop between them.
    """
    SerializedTypeID = 0

    @abc.abstractmethod
    def to_serializable(self):
        raise NotImplementedError()

    @classmethod
    @abc.abstractmethod
    def from_serializable(cls, data):
        raise NotImplementedError()


class SerializableItemManager:
    SERIALIZED_MIMETYPE = "/serialized-objects/data/"
    SERIALIZED_OBJECT_ID_MIMETYPE = "/serialized-objects/type-id/"

    _typeID_itemType_map = defaultdict(lambda: dict())

    @classmethod
    def addItemsToMimedata(cls, items, mimedata: QtCore.QMimeData):
        typeIds = [item.SerializedTypeID for item in items if isinstance(item, SerializableItemMixin)]
        mimedata.setData(cls.SERIALIZED_OBJECT_ID_MIMETYPE, cls.serializeData(typeIds))
        mimedata.setData(cls.SERIALIZED_MIMETYPE, cls.serializeItems(items))

    @classmethod
    def getItemsFromMimedata(cls, caller, mimedata: QtCore.QMimeData):
        if not (mimedata.hasFormat(cls.SERIALIZED_MIMETYPE) and mimedata.hasFormat(cls.SERIALIZED_OBJECT_ID_MIMETYPE)):
            return []

        typeIds = cls.deserializeData(mimedata.data(cls.SERIALIZED_OBJECT_ID_MIMETYPE))
        object_data = cls.deserializeData(mimedata.data(cls.SERIALIZED_MIMETYPE))

        return cls.deserializeItemsFromSerializable(caller, typeIds, object_data)


    @classmethod
    def serializeItems(cls, items: Iterable[object]):
        serializable_data = [item.to_serializable() for item in items if isinstance(item, SerializableItemMixin)]
        return cls.serializeData(serializable_data)

    @classmethod
    def serializeData(cls, serializable_data):
        return QtCore.QByteArray(json.dumps(serializable_data).encode())

    @classmethod
    def deserializeData(cls, serialized: QtCore.QByteArray):
        return json.loads(str(serialized.data(), encoding="utf-8"))


    @classmethod
    def deserializeItemsFromSerializable(cls, caller, typeIds, data):
        items = []
        for typeId, serialized in zip(typeIds, data):
            typeKey = cls.getTypeKey(caller)
            item_type: SerializableItemMixin = cls._typeID_itemType_map[typeId].get(typeKey)
            if not item_type: continue
            items.append(item_type.from_serializable(serialized))
        return items


    @classmethod
    def getTypeKey(cls, caller):
        if not inspect.isclass(caller):
            caller = caller.__class__

        if issubclass(caller, (QtWidgets.QTreeWidget, QtWidgets.QTreeWidgetItem)):
            return "tree"
        if issubclass(caller, (QtWidgets.QTableWidget, QtWidgets.QTableWidgetItem)):
            return "table"
        if issubclass(caller, (QtWidgets.QListWidget, QtWidgets.QListWidgetItem)):
            return "list"
        return None

    @classmethod
    def registerCustomWidgetItemType(cls, object_type: Type[SerializableItemMixin]):
        typeKey = cls.getTypeKey(object_type)
        cls._typeID_itemType_map[object_type.SerializedTypeID][typeKey] = object_type


class SerializableListWidget(QtWidgets.QListWidget):
    def mimeData(self, items):
        data = QtWidgets.QListWidget.mimeData(self, items)
        SerializableItemManager.addItemsToMimedata(items, data)
        return data

    def dropMimeData(self, index:int, data:QtCore.QMimeData, action:QtCore.Qt.DropAction) -> bool:
        items = SerializableItemManager.getItemsFromMimedata(self, data)
        if items:
            for i, item in enumerate(items):
                self.insertItem(index+i, item)
            return True
        else:
            return QtWidgets.QListWidget.dropMimeData(self, index, data, action)

class SerializableTreeWidget(QtWidgets.QTreeWidget):
    def mimeData(self, items):
        data = QtWidgets.QTreeWidget.mimeData(self, items)
        SerializableItemManager.addItemsToMimedata(items, data)
        return data

    def dropMimeData(self, parent:QtWidgets.QTreeWidgetItem, index:int, data:QtCore.QMimeData, action:QtCore.Qt.DropAction) -> bool:
        items = SerializableItemManager.getItemsFromMimedata(self, data)
        if items:
            parent = self.invisibleRootItem() if parent is None else parent
            parent.setExpanded(True)
            for i, item in enumerate(items):
                parent.insertChild(index+i, item)
            return True
        else:
            return QtWidgets.QTreeWidget.dropMimeData(self, parent, index, data, action)

class SerializableTableWidget(QtWidgets.QTableWidget):
    def mimeData(self, items):
        if not items: return None
        data = QtWidgets.QTableWidget.mimeData(self, items)
        SerializableItemManager.addItemsToMimedata(items, data)
        return data

    def dropMimeData(self, row:int, column:int, data:QtCore.QMimeData, action:QtCore.Qt.DropAction) -> bool:
        items = SerializableItemManager.getItemsFromMimedata(self, data)
        if items:
            self.setItem(row, column, items[0])
            return True
        else:
            return QtWidgets.QTableWidget.dropMimeData(self, row, column, data, action)


def make_widget_serializable(widget):
    """
    Monkeypatches an existing widget to add handling for the SerializableItemMixin
    in its drag and drop functionality. It will overwrite the default mimeData functions.
    """
    if isinstance(widget, QtWidgets.QListWidget):
        widget.mimeData = partial(SerializableListWidget.mimeData, widget)
        widget.dropMimeData = partial(SerializableListWidget.dropMimeData, widget)
    elif isinstance(widget, QtWidgets.QTreeWidget):
        widget.mimeData = partial(SerializableTreeWidget.mimeData, widget)
        widget.dropMimeData = partial(SerializableTreeWidget.dropMimeData, widget)
    elif isinstance(widget, QtWidgets.QTableWidget):
        widget.mimeData = partial(SerializableTableWidget.mimeData, widget)
        widget.dropMimeData = partial(SerializableTableWidget.dropMimeData, widget)




# ----------------------------------------------------------
# Example

class MyDataType:
    def __init__(self, name, age):
        self.name = name
        self.age = age

    @classmethod
    def from_dictionary(cls, data):
        obj = cls(None, None)
        obj.__dict__ = data
        return obj

    def to_dictionary(self):
        return self.__dict__



class SerializableMyDataTypeMixin(SerializableItemMixin):
    def __init__(self, data):
        super(SerializableItemMixin, self).__init__(data)
        self.myData = data

    def to_serializable(self):
        return self.myData.to_dictionary()

    @classmethod
    def from_serializable(cls, data):
        return cls(MyDataType.from_dictionary(data))


class MyDataListWidgetItem(QtWidgets.QListWidgetItem, SerializableMyDataTypeMixin):
    def __init__(self, data):
        super(MyDataListWidgetItem, self).__init__()
        self.myData = data
        self.setText(data.name + " - " + str(data.age))
        self.setFlags(self.flags() | QtCore.Qt.ItemIsDragEnabled)

SerializableItemManager.registerCustomWidgetItemType(MyDataListWidgetItem)



class MyDataTreeWidgetItem(QtWidgets.QTreeWidgetItem, SerializableMyDataTypeMixin):
    def __init__(self, data):
        super(MyDataTreeWidgetItem, self).__init__()
        self.myData = data
        self.setText(0, data.name)
        self.setText(1, str(data.age))
        self.setFlags(self.flags() | QtCore.Qt.ItemIsDragEnabled)

SerializableItemManager.registerCustomWidgetItemType(MyDataTreeWidgetItem)



class MyDataTableWidgetItem(QtWidgets.QTableWidgetItem, SerializableMyDataTypeMixin):
    def __init__(self, data):
        super(MyDataTableWidgetItem, self).__init__()
        self.myData = data
        self.setText(data.name + ", " + str(data.age))
        self.setFlags(self.flags() | QtCore.Qt.ItemIsDragEnabled)

SerializableItemManager.registerCustomWidgetItemType(MyDataTableWidgetItem)



class TestDialog(QtWidgets.QDialog):
    def __init__(self, a_data, b_data):
        super(TestDialog, self).__init__()
        self.a_data = a_data
        self.b_data = b_data

        self.__setupUI()
        self.populateLists()
        self.__setupCallbacks()

    def __setupUI(self):
        layout = QtWidgets.QVBoxLayout()
        self.setLayout(layout)

        self.listWidgetA = QtWidgets.QListWidget()

        self.listWidgetA.setDragEnabled(True)
        self.listWidgetA.setAcceptDrops(False)
        self.listWidgetA.setDragDropMode(QtWidgets.QAbstractItemView.DragOnly)
        layout.addWidget(self.listWidgetA)

        self.listWidgetB = QtWidgets.QListWidget()
        self.listWidgetB.setAcceptDrops(True)
        self.listWidgetB.setDragEnabled(True)
        self.listWidgetB.setDragDropMode(QtWidgets.QAbstractItemView.DragDrop)
        self.listWidgetB.setDefaultDropAction(QtCore.Qt.MoveAction)
        layout.addWidget(self.listWidgetB)

        self.treeWidgetB = QtWidgets.QTreeWidget()
        self.treeWidgetB.setColumnCount(2)
        self.treeWidgetB.setHeaderLabels(["name", "age"])
        self.treeWidgetB.setAcceptDrops(True)
        self.treeWidgetB.setDragEnabled(True)
        self.treeWidgetB.setDragDropMode(QtWidgets.QAbstractItemView.DragDrop)
        self.treeWidgetB.setDefaultDropAction(QtCore.Qt.MoveAction)
        layout.addWidget(self.treeWidgetB)

        self.tableWidgetB = QtWidgets.QTableWidget()
        self.tableWidgetB.setRowCount(4)
        self.tableWidgetB.setColumnCount(4)
        self.tableWidgetB.setAcceptDrops(True)
        self.tableWidgetB.setDragEnabled(True)
        self.tableWidgetB.setDragDropMode(QtWidgets.QAbstractItemView.DragDrop)
        self.tableWidgetB.setDefaultDropAction(QtCore.Qt.MoveAction)
        layout.addWidget(self.tableWidgetB)

        self.testButton = QtWidgets.QPushButton("test!")
        layout.addWidget(self.testButton)

        self.setWindowFlags(QtCore.Qt.WindowStaysOnTopHint)

        make_widget_serializable(self.listWidgetA)
        make_widget_serializable(self.listWidgetB)
        make_widget_serializable(self.treeWidgetB)
        make_widget_serializable(self.tableWidgetB)

    def __setupCallbacks(self):
        self.testButton.clicked.connect(self.runTest)

    def populateLists(self):
        self.listWidgetA.clear()
        self.listWidgetB.clear()

        for data in self.a_data:
            item = MyDataListWidgetItem(data)
            self.listWidgetA.addItem(item)
        self.listWidgetA.addItem(QtWidgets.QListWidgetItem("special case"))

        for data in self.b_data:
            item = MyDataListWidgetItem(data)
            self.listWidgetB.addItem(item)

    def runTest(self):
        print("testing!")
        for i in range(self.listWidgetB.count()):
            item = self.listWidgetB.item(i)
            print(item)



if __name__ == '__main__':
    a_data = [
        MyDataType("Mitchell", 24),
        MyDataType("Jason", 35),
        MyDataType("Bourne", 44),
    ]
    b_data = [
        MyDataType("Watney", 29),
    ]
    w = TestDialog(a_data, b_data)
    w.show()