PyQt Drag and drop: Outliner-like QTreeView in Maya

Ok, it's definitely been a while since I last posted. I promise I'm not dead, just very busy with work! Anyways, this was an issue that came up for me a while ago -- one of those things I had to spend a couple days puzzling over. The issue was implementing internal drag and drop in a PyQt hierarchical tree display using the model/view framework.

Enabling drag and drop for any of PyQt's convenience view classes (QTreeWidget, QListWidget) is fairly straightforward (alright, it's downright easy). Since you have to pretty much implement it from scratch when you use QAbstractItemModel paired with QTreeView, there end up being a couple of 'gotchas'.

If you aren't familiar with the model/view paradigm in PyQt, I'd like to point you towards Yasin Uludag's excellent video series here -- they should give you a pretty good introduction. If you've seen his videos before, the code in my examples should appear fairly familiar, as I've structured things pretty much the same.

By the end of this post, you should be able to implement your own outliner-like window in Maya using PyQt, as well as understand a little about how to integrate the Qt window with objects in the Maya scene using PyMEL.

Let's jump in:

from pymel.core import *
from PyQt4 import QtCore, QtGui
import cPickle

import maya.OpenMayaUI
import sip

def mayaMainWindow():
    '''Returns the Maya window as a QMainWindow instance.'''
    ptr = maya.OpenMayaUI.MQtUtil.mainWindow()
    if ptr is not None:
        return sip.wrapinstance( long(ptr), QtCore.QObject )

We'll start out with some basic imports. PyMEL will be used for interacting with Maya (although you could also use maya.cmds if you absolutely had to), PyQt4 for drawing the GUI elements, cPickle or pickle for serializing data. The OpenMaya and sip imports are for the first function, the usual boilerplate Maya Main Window function that allows us to make the windows we create children of the Maya window.

class TreeItem( object ):
    '''Wraps a PyNode to provide an interface to the QAbstractItemModel'''
    
    def __init__( self, node=None ):
        '''
        Instantiates a new tree item wrapping a PyNode DAG object.
        '''
        self.node = node
        
        self.parent = None
        self.children = []
        
    @property
    def displayName( self ):
        '''Returns the wrapped PyNode's name.'''
        return self.node.name()
    @displayName.setter
    def displayName( self, name ):
        '''Renames the wrapped PyNode.'''
        self.node.rename( str( name ) )
    
    def addChild( self, child ):
        '''Adds a given item as a child of this item.
        
        Also handles parenting of the PyNode in the maya scene
        '''
        self.children.append( child )
        child.parent = self
        
        # If adding a child to the root, parent the node to the world
        if self.node is None:
            # In earlier versions of PyMEL, parenting an object
            # to its current parent throws an error
            if child.node.getParent() is not None:
                child.node.setParent( world=True )
                return
                
        if child.node.getParent() != self.node:
            child.node.setParent( self.node )
        
    def numChildren( self ):
        '''Returns the number of child items.'''
        return len( self.children )
        
    def removeChildAtRow( self, row ):
        '''Removes an item at the given index from the list of children.'''
        self.children.pop( row )
    
    def childAtRow( self, row ):
        '''Retrieves the item at the given index from the list of children.'''
        return self.children[row]
        
    def row( self ):
        '''Get this item's index in its parent item's child list.'''
        if self.parent:
            return self.parent.children.index( self )
        return 0
            
    def log( self, level=-1 ):
        '''Returns a textual representation of an item's hierarchy.'''
        level += 1
        
        output = ''
        for i in range( level ):
            output += '\t'
        
        output += self.node.name() if self.node is not None else 'Root'
        output += '\n'
        
        for child in self.children:
            output += child.log( level )
        
        level -= 1
        
        return output

Here's the first class. Since we will use Qt's model/view interface to create the tree view, we need an interface to the Maya objects that will be represented in the outliner. This 'item' class can wrap a PyNode to provide a usable interface to the object, returning information regarding the object's name and hierarchy; it can also edit the PyNode, so that the outliner can affect changes in the Maya scene as well.

For this example, the PyNodes in question will represent joints, since they have a handy visual representation of their hierarchy - it will be easy to see hierarchical relationships between joints as they change as a result of dragging and dropping items in the outliner.

def gatherItems():
    '''Return a scene hierarchy of top-level joints as TreeItems.
    
    Creates and returns a root TreeItem with items for all joints that are
    direct children of the world as children.
    '''
    # Create a null TreeItem to serve as the root
    rootItem = TreeItem()
    
    topLevelJoints = [jnt for jnt in ls( type='joint' ) if jnt.getParent() is None]
    
    def recursiveCreateItems( node ):
        # An inline function for recursively creating tree items and adding them to their parent
        item = TreeItem( node )
        for child in node.getChildren( type='joint' ):
            childItem = recursiveCreateItems( child )
            item.addChild( childItem )
        return item
    
    for jnt in topLevelJoints:
        item = recursiveCreateItems( jnt )
        rootItem.addChild( item )
    
    return rootItem

This function gathers the data from the scene and compiles the hierarchy of TreeItems. It recursively traverses through all the top-level joints in the scene (joints with no parent), wraps them in a TreeItem instance, and adds that item to the 'children' list of the appropriate parent. It then returns a null TreeItem that acts as a 'world' level item, with all those top-level joints as its children.

class OutlinerModel( QtCore.QAbstractItemModel ):
    '''A drag and drop enabled, editable, hierarchical item model.'''
    
    def __init__( self, root ):
        '''Instantiates the model with a root item.'''
        super( OutlinerModel, self ).__init__()
        self.root = root
        
    def itemFromIndex( self, index ):
        '''Returns the TreeItem instance from a QModelIndex.'''
        return index.internalPointer() if index.isValid() else self.root
        
    def rowCount( self, index ):
        '''Returns the number of children for the given QModelIndex.'''
        item = self.itemFromIndex( index )
        return item.numChildren()
    
    def columnCount( self, index ):
        '''This model will have only one column.'''
        return 1
    
    def flags( self, index ):
        '''Valid items are selectable, editable, and drag and drop enabled. Invalid indices (open space in the view)
        are also drop enabled, so you can drop items onto the top level.
        '''
        if not index.isValid():
            return QtCore.Qt.ItemIsEnabled | QtCore.Qt.ItemIsDropEnabled
        
        return QtCore.Qt.ItemIsEnabled | QtCore.Qt.ItemIsDropEnabled | QtCore.Qt.ItemIsDragEnabled |\
               QtCore.Qt.ItemIsSelectable | QtCore.Qt.ItemIsEditable
       
    def supportedDropActions( self ):
        '''Items can be moved and copied (but we only provide an interface for moving items in this example.'''
        return QtCore.Qt.MoveAction | QtCore.Qt.CopyAction
    
    def headerData( self, section, orientation, role ):
        '''Return the header title.'''
        if section == 0 and orientation == QtCore.Qt.Horizontal:
            if role == QtCore.Qt.DisplayRole:
                return 'Joints'
        return QtCore.QVariant()
    
    def data( self, index, role ):
        '''Return the display name of the PyNode from the item at the given index.'''
        if role == QtCore.Qt.DisplayRole or role == QtCore.Qt.EditRole:
            item = self.itemFromIndex( index )
            return item.displayName
            
    def setData( self, index, value, role ):
        '''Set the name of the PyNode from the item being edited.'''
        item = self.itemFromIndex( index )
        item.displayName = str( value.toString() )
        self.dataChanged.emit( QtCore.QModelIndex(), QtCore.QModelIndex() )
        return True
        
    def index( self, row, column, parentIndex ):
        '''Creates a QModelIndex for the given row, column, and parent.'''
        if not self.hasIndex( row, column, parentIndex ):
            return QtCore.QModelIndex()
            
        parent = self.itemFromIndex( parentIndex )
        return self.createIndex( row, column, parent.childAtRow( row ) )
    
    def parent( self, index ):
        '''Returns a QMoelIndex for the parent of the item at the given index.'''
        item = self.itemFromIndex( index )
        parent = item.parent
        if parent == self.root:
            return QtCore.QModelIndex()
        return self.createIndex( parent.row(), 0, parent )
    
    def insertRows( self, row, count, parentIndex ):
        '''Add a number of rows to the model at the given row and parent.'''
        self.beginInsertRows( parentIndex, row, row+count-1 )
        self.endInsertRows()
        return True
    
    def removeRows( self, row, count, parentIndex ):
        '''Remove a number of rows from the model at the given row and parent.'''
        self.beginRemoveRows( parentIndex, row, row+count-1 )
        parent = self.itemFromIndex( parentIndex )
        for x in range( count ):
            parent.removeChildAtRow( row )
        self.endRemoveRows()
        return True
    
    def mimeTypes( self ):
        '''The MimeType for the encoded data.'''
        types = QtCore.QStringList( 'application/x-pynode-item-instance' )
        return types
    
    def mimeData( self, indices ):
        '''Encode serialized data from the item at the given index into a QMimeData object.'''
        data = ''
        item = self.itemFromIndex( indices[0] )
        try:
            data += cPickle.dumps( item )
        except:
            pass
        mimedata = QtCore.QMimeData()
        mimedata.setData( 'application/x-pynode-item-instance', data )
        return mimedata
    
    def dropMimeData( self, mimedata, action, row, column, parentIndex ):
        '''Handles the dropping of an item onto the model.
        
        De-serializes the data into a TreeItem instance and inserts it into the model.
        '''
        if not mimedata.hasFormat( 'application/x-pynode-item-instance' ):
            return False
        item = cPickle.loads( str( mimedata.data( 'application/x-pynode-item-instance' ) ) )
        dropParent = self.itemFromIndex( parentIndex )
        dropParent.addChild( item )
        self.insertRows( dropParent.numChildren()-1, 1, parentIndex )
        self.dataChanged.emit( parentIndex, parentIndex )
        return True

Here is the data model that will effectively translate the collection of data (TreeItems) into a form that can be displayed in our outliner (a QTreeView). The model is responsible for acting as a middle man between the GUI object and the data.

There are a few important things to note here. The 'flags' function (line 137) declares that items can not only be selected and edited, but also dragged and dropped. Line 173 in the 'index' function is also very important - sometimes the data model will request information about indices that don't exist, leading to 'list index out of range' errors. I had a hell of a time figuring this out, and I still don't fully understand why this happens. However, looking through the C++ Qt source code, I found that convenience classes like QStandardItemModel, which inherits QAbstractItemModel does in fact employ a similar logical check to avoid queries to invalid indices.

Perhaps the most important bit to call attention to are the last three functions, 'mimeTypes', 'mimeData', and 'dropMimeData'. The first declares acceptable Mime types for the model. Since we aren't using any standard data type, what we put here is fairly arbitrary. Check out good old wikipedia for more info on Mime types.

class SkeletonOutliner( QtGui.QMainWindow ):
    '''A window containing a tree view set up for drag and drop.'''
    
    def __init__( self, parent=mayaMainWindow() ):
        '''Instantiates the window as a child of the Maya main window, sets up the
        QTreeView with an OutlinerModel, and enables the drag and drop operations.
        '''
        super( SkeletonOutliner, self ).__init__( parent )
        
        self.tree = QtGui.QTreeView()
        self.outlinerModel = OutlinerModel( gatherItems() )
        self.tree.setModel( self.outlinerModel )
        self.tree.setDragEnabled( True )
        self.tree.setAcceptDrops( True )
        self.tree.setDragDropMode( QtGui.QAbstractItemView.InternalMove )
        
        self.selModel = self.tree.selectionModel()
        self.selModel.currentChanged.connect( self.selectInScene )
        
        self.tree.expandAll()
        
        self.setCentralWidget( self.tree )
        
        self.show()
        
    def selectInScene( self, current, previous ):
        '''Callback for selecting the PyNode in the maya scene when the outliner selection changes.'''
        pynode = self.outlinerModel.itemFromIndex( current ).node
        select( pynode, r=True )

The last part of the code defines a Qt window, parented to Maya's main window with a tree view hooked up to our OutlinerModel. There's a callback hooked up to the tree view's selection model that will change the selection in the Maya scene to the item selected in the outliner. And that's really about it.

Call the window like this:

outliner = SkeletonOutliner()

To see it in action, create a new scene in Maya and create a few joint chains, then instantiate the SkeletonOutliner class. You should be able to select joints, drag and drop them to reparent, and double click the names to rename them.