More work on NBT editor

Adding to Compounds is possible, showing a tag type menu.
Adding to an empty list uses the menu too
setModel changed to setRootTag
tries to keep tree expand state when model changes (e.g. due to revision change)

chunk editor treeView->nbtEditor

src/mcedit2/editortools/edit_chunk.py

@@ -35,6 +35,8 @@ class ChunkTool(EditorTool):
         self.toolWidget.cxSpinBox.valueChanged.connect(self.chunkPositionDidChange)
         self.toolWidget.czSpinBox.valueChanged.connect(self.chunkPositionDidChange)
 
+        self.toolWidget.nbtEditor.editorSession = self.editorSession
+
         self.currentChunk = None
         self.selectionNode = None
         self.overlayNode = scenegraph.Node()
@@ -126,15 +128,10 @@ class ChunkTool(EditorTool):
     def updateNBTView(self):
         chunk = self.currentChunk
         if chunk is None:
-            self.toolWidget.nbtEditor.setModel(None)
+            self.toolWidget.nbtEditor.setRootTag(None)
             return
 
-        model = NBTTreeModel(chunk.rootTag)
+        self.toolWidget.nbtEditor.setRootTag(chunk.rootTag)
-
-        self.toolWidget.nbtEditor.setModel(model)
-        # self.toolWidget.nbtEditor.expandToDepth(0)
-        # self.toolWidget.nbtEditor.resizeColumnToContents(0)
-        # self.toolWidget.nbtEditor.resizeColumnToContents(1)
 
         self.toolWidget.cxSpinBox.setValue(chunk.cx)
         self.toolWidget.czSpinBox.setValue(chunk.cz)

src/mcedit2/editortools/edit_entity.py

@@ -47,6 +47,7 @@ class EntityTool(EditorTool):
     def createToolWidget(self):
         self.toolWidget = load_ui("editortools/select_entity.ui")
         self.toolWidget.entityListBox.currentIndexChanged.connect(self.setSelectedEntity)
+        self.toolWidget.nbtEditor.editorSession = self.editorSession
 
     def mousePress(self, event):
         command = SelectEntityCommand(self, event.ray)
@@ -69,10 +70,9 @@ class EntityTool(EditorTool):
 
     def setSelectedEntity(self, index):
         if len(self.selectedEntities):
-            model = NBTTreeModel(self.selectedEntities[index].raw_tag())
+            self.toolWidget.nbtEditor.setRootTag(self.selectedEntities[index].raw_tag())
-            self.toolWidget.treeView.setModel(model)
         else:
-            self.toolWidget.treeView.setModel(None)
+            self.toolWidget.nbtEditor.setRootTag(None)
 
 
 def entitiesOnRay(dimension, ray, rayWidth=2.0, maxDistance = 1000):

src/mcedit2/panels/player.py

@@ -22,9 +22,6 @@ log = logging.getLogger(__name__)
 class PlayerPropertyChangeCommand(SimpleRevisionCommand):
     pass
 
-class NBTDataChangeCommand(SimpleRevisionCommand):
-    pass
-
 class PlayerPanel(QtGui.QWidget):
     def __init__(self, editorSession, *args, **kwargs):
         """
@@ -72,6 +69,9 @@ class PlayerPanel(QtGui.QWidget):
         self.editorSession.revisionChanged.connect(self.revisionDidChange)
         self.initPropertiesWidget()
 
+        self.nbtEditor.editorSession = self.editorSession
+        self.nbtEditor.editMade.connect(self.nbtEditWasMade)
+
         centerWidgetInScreen(self)
 
     def initPropertiesWidget(self):
@@ -109,34 +109,16 @@ class PlayerPanel(QtGui.QWidget):
         model.propertyChanged.connect(self.propertyDidChange)
 
     def updateNBTTree(self):
-        model = NBTTreeModel(self.selectedPlayer.rootTag)
+        self.nbtEditor.undoCommandPrefixText = ("Player %s: " % self.selectedUUID) if self.selectedUUID else "Single-player: "
-        model.dataChanged.connect(self.nbtDataDidChange)
+        self.nbtEditor.setRootTag(self.selectedPlayer.rootTag)
-        self.nbtEditor.setModel(model)
+
+    def nbtEditWasMade(self):
+        self.selectedPlayer.dirty = True
 
     def revisionDidChange(self):
         self.initPropertiesWidget()
         self.updateNBTTree()
 
-    def nbtDataDidChange(self, index):
-        model = self.nbtEditor.model
-        parent = model.parent(index)
-        item = model.getItem(index)
-        if parent is not None and parent.isList:
-            name = str(parent.tag.index(item.tag))
-        else:
-            name = item.tag.name
-
-        if self.selectedUUID != "":
-            text = "Change player %s NBT tag %s" % (self.selectedUUID, name)
-        else:
-            text = "Change single-player NBT tag %s" % name
-
-        command = NBTDataChangeCommand(self.editorSession, text)
-        with command.begin():
-            self.selectedPlayer.dirty = True
-            self.editorSession.worldEditor.syncToDisk()
-        self.editorSession.pushCommand(command)
-
     def propertyDidChange(self, name, value):
         if self.selectedUUID != "":
             text = "Change player %s property %s" % (self.selectedUUID, name)
@@ -169,7 +151,7 @@ class PlayerPanel(QtGui.QWidget):
             self.selectedUUID = UUID
         except PlayerNotFound:
             log.info("PlayerPanel: player %s not found!", UUID)
-            self.nbtEditor.setModel(None)
+            self.nbtEditor.setRootTag(None)
         else:
             self.updateNBTTree()
 

src/mcedit2/ui/editortools/select_entity.ui

@@ -29,17 +29,25 @@
     </layout>
    </item>
    <item>
-    <widget class="QTreeView" name="treeView">
+    <widget class="NBTEditorWidget" name="nbtEditor" native="true">
-     <property name="alternatingRowColors">
+     <property name="sizePolicy">
-      <bool>true</bool>
+      <sizepolicy hsizetype="Preferred" vsizetype="Preferred">
-     </property>
+       <horstretch>0</horstretch>
-     <property name="sortingEnabled">
+       <verstretch>1</verstretch>
-      <bool>true</bool>
+      </sizepolicy>
      </property>
     </widget>
    </item>
   </layout>
  </widget>
+ <customwidgets>
+  <customwidget>
+   <class>NBTEditorWidget</class>
+   <extends>QWidget</extends>
+   <header>nbteditorwidget.h</header>
+   <container>1</container>
+  </customwidget>
+ </customwidgets>
  <resources/>
  <connections/>
 </ui>

src/mcedit2/widgets/nbttree/nbttreemodel.py

@@ -28,6 +28,7 @@ _iconTypes = [
     "array.png", # 12 - shortarray
 ]
 
+NBTPathRole = QtCore.Qt.UserRole + 1
 
 def NBTIcon(type):
     icon = _nbtIcons.get(type)
@@ -98,15 +99,21 @@ class NBTTreeCompound(object):
             return "%s items" % len(tag)
         return summary
 
-    def insertChildren(self, position, count):
+    def insertChildren(self, position, count, tagID):
         if position < 0 or position > len(self.childItems):
             return False
 
         for row in range(count):
-            data = nbt.TAG_Byte()
+            name = "Unnamed"
-            self.tag.insert(position + row, data)
+            i = 0
+            while name in self.tag:
+                i += 1
+                name = "Unnamed %d" % i
 
-            item = NBTTreeItem(data, self)
+            tag = nbt.tag_classes[tagID]()
+            self.tag[name] = tag
+
+            item = NBTTreeItem(tag, self)
             self.childItems.insert(position + row, item)
 
         return True
@@ -119,13 +126,23 @@ class NBTTreeCompound(object):
             return False
 
         for row in range(count):
-            self.childItems.pop(position)
+            name = self.childItems.pop(position).tag.name
+            del self.tag[name]
 
         return True
 
     def setValue(self, value):
         return False
 
+    def nbtPath(self, child=None):
+        if self.parentItem is None:
+            path = []
+        else:
+            path = self.parentItem.nbtPath(self)
+        if child:
+            path.append(child.tag.name)
+        return path
+
 
 class NBTTreeList(object):
     isCompound = False
@@ -163,14 +180,14 @@ class NBTTreeList(object):
 
             return ", ".join((fmt % i.value) for i in self.tag)
 
-    def insertChildren(self, position, count):
+    def insertChildren(self, position, count, tagID):
         if position < 0 or position > len(self.childItems):
             return False
 
         for row in range(count):
-            data = nbt.tag_classes[self.tag.list_type or nbt.ID_BYTE]()
+            tag = nbt.tag_classes[self.tag.list_type or tagID]()
-            self.tag.insert(position + row, data)
+            self.tag.insert(position + row, tag)
-            item = NBTTreeItem(data, self)
+            item = NBTTreeItem(tag, self)
             self.childItems.insert(position + row, item)
 
         return True
@@ -184,12 +201,23 @@ class NBTTreeList(object):
 
         for row in range(count):
             self.childItems.pop(position)
+            self.tag.pop(position)
 
         return True
 
     def setValue(self, value):
         return False
 
+    def nbtPath(self, child=None):
+        if self.parentItem is None:
+            path = []
+        else:
+            path = self.parentItem.nbtPath(self)
+        if child:
+            row = self.childItems.index(child)
+            path.append(row)
+
+        return path
 
 class NBTTreeItem(object):
     isCompound = False
@@ -229,9 +257,15 @@ class NBTTreeItem(object):
         self.tag.value = value
         return True
 
+    def nbtPath(self):
+        if self.parentItem is None:
+            return []
+        return self.parentItem.nbtPath(self)
+
 class NBTTreeModel(QtCore.QAbstractItemModel):
     def __init__(self, rootTag, parent=None):
         super(NBTTreeModel, self).__init__(parent)
+        self._internalPointers = {}
 
         self.rootItem = MakeNBTTreeItem(rootTag, None)
         self.rootTag = rootTag
@@ -245,20 +279,13 @@ class NBTTreeModel(QtCore.QAbstractItemModel):
     def tagID(self, index):
         return self.getItem(index).tag.tagID
 
-    def data(self, index, role=QtCore.Qt.DisplayRole):
+    # --- Data ---
-        item = self.getItem(index)
-        column = index.column()
 
-        if role == QtCore.Qt.DecorationRole:
+    def headerData(self, section, orientation, role=QtCore.Qt.DisplayRole):
-            if column == 0:
+        if orientation == QtCore.Qt.Horizontal and role == QtCore.Qt.DisplayRole:
-                return NBTIcon(item.tag.tagID)
+            return ("Name", "Value", "", "")[section]
-            if column == 2:
-                return self.addIcon if item.isList or item.isCompound else None
-            if column == 3:
-                return self.removeIcon
 
-        if role in (QtCore.Qt.DisplayRole, QtCore.Qt.EditRole):
+        return None
-            return item.data(column)
 
     def flags(self, index):
         flags = QtCore.Qt.ItemIsEnabled | QtCore.Qt.ItemIsSelectable
@@ -269,55 +296,85 @@ class NBTTreeModel(QtCore.QAbstractItemModel):
             flags |= QtCore.Qt.ItemIsEditable
         return flags
 
+    def data(self, index, role=QtCore.Qt.DisplayRole):
+        item = self.getItem(index)
+        column = index.column()
+
+        if role == QtCore.Qt.DecorationRole:
+            if column == 0:
+                return NBTIcon(item.tag.tagID)
+            if column == 2:
+                return self.addIcon if item.isList or item.isCompound else None
+            if column == 3:
+                return self.removeIcon if item is not self.rootItem else None
+
+        if role in (QtCore.Qt.DisplayRole, QtCore.Qt.EditRole):
+            return item.data(column)
+
+        if role == NBTPathRole:
+            return item.nbtPath()
+
+    # --- Structure ---
 
     def getItem(self, index):
         if index.isValid():
-            item = index.internalPointer()
+            item = self._internalPointers[index.internalId()]
             if item:
                 return item
         else:
-            return self.rootItem
+            return None
 
-    def headerData(self, section, orientation, role=QtCore.Qt.DisplayRole):
+    def rowCount(self, parent=QtCore.QModelIndex()):
-        if orientation == QtCore.Qt.Horizontal and role == QtCore.Qt.DisplayRole:
+        if not parent.isValid():
-            return ("Name", "Value", "", "")[section]
+            return 1
+        parentItem = self.getItem(parent)
 
-        return None
+        return parentItem.childCount()
+
+    def createIndex(self, row, column, item):
+        self._internalPointers[id(item)] = item
+        return super(NBTTreeModel, self).createIndex(row, column, id(item))
 
     def index(self, row, column, parent=QtCore.QModelIndex()):
         if not parent.isValid():
+            assert row == 0
             return self.createIndex(row, column, self.rootItem)
 
         parentItem = self.getItem(parent)
-        childItem = parentItem.child(row)
+        if parentItem is None:
-        if childItem:
-            return self.createIndex(row, column, childItem)
-        else:
             return QtCore.QModelIndex()
 
-    def insertRow(self, position, parent=QtCore.QModelIndex()):
+        childItem = parentItem.child(row)
-        return self.insertRows(position, 1, parent)
+        if childItem is None:
+            return QtCore.QModelIndex()
 
-    def insertRows(self, row, count, parent=QtCore.QModelIndex()):
+        return self.createIndex(row, column, childItem)
-        parentItem = self.getItem(parent)
-        self.beginInsertRows(parent, row, row + count - 1)
-        success = parentItem.insertChildren(row, count)
-        self.endInsertRows()
-
-        return success
 
     def parent(self, index):
         if not index.isValid():
             return QtCore.QModelIndex()
 
-        childItem = self.getItem(index)
+        item = self.getItem(index)
-        parentItem = childItem.parent()
+        parentItem = item.parent()
 
-        if parentItem is None:
+        if parentItem is None:  # item is self.rootItem
             return QtCore.QModelIndex()
 
         return self.createIndex(parentItem.childNumber(), 0, parentItem)
 
+    # --- Editing ---
+
+    def insertRow(self, position, parent=QtCore.QModelIndex(), tagID=None):
+        return self.insertRows(position, 1, parent, tagID)
+
+    def insertRows(self, row, count, parent=QtCore.QModelIndex(), tagID=None):
+        parentItem = self.getItem(parent)
+        self.beginInsertRows(parent, row, row + count - 1)
+        success = parentItem.insertChildren(row, count, tagID)
+        self.endInsertRows()
+
+        return success
+
     def removeRow(self, position, parent=QtCore.QModelIndex()):
         self.removeRows(position, 1, parent)
 
@@ -330,13 +387,6 @@ class NBTTreeModel(QtCore.QAbstractItemModel):
 
         return success
 
-    def rowCount(self, parent=QtCore.QModelIndex()):
-        if not parent.isValid():
-            return 1
-        parentItem = self.getItem(parent)
-
-        return parentItem.childCount()
-
     def setData(self, index, value, role=QtCore.Qt.EditRole):
         if role != QtCore.Qt.EditRole:
             return False

src/mcedit2/widgets/nbttree/nbttreeview.py

@@ -6,17 +6,28 @@ import logging
 
 from PySide import QtGui, QtCore
 from PySide.QtCore import Qt
-from mcedit2.util.resources import resourcePath
+from mcedit2.command import SimpleRevisionCommand
 
-from mcedit2.widgets.nbttree.nbttreemodel import NBTFilterProxyModel
+from mcedit2.widgets.nbttree.nbttreemodel import NBTFilterProxyModel, NBTPathRole, NBTIcon, NBTTreeModel
 from mcedit2.util.load_ui import registerCustomWidget
-from mcedit2.widgets.layout import Row, Column
+from mcedit2.widgets.layout import Column
 
 
 log = logging.getLogger(__name__)
 
+
+class NBTDataChangeCommand(SimpleRevisionCommand):
+    pass
+
 @registerCustomWidget
 class NBTEditorWidget(QtGui.QWidget):
+    undoCommandPrefixText = ""
+    editorSession = NotImplemented
+    proxyModel = None
+    rootTag = None
+
+    editMade = QtCore.Signal()  # emitted to allow clients to mark the NBT tree's parent structure as dirty - xxx really??
+
     def __init__(self, *args, **kwargs):
         super(NBTEditorWidget, self).__init__(*args, **kwargs)
         self.model = None
@@ -27,37 +38,197 @@ class NBTEditorWidget(QtGui.QWidget):
 
         self.setLayout(Column(self.treeView))
 
-    def setModel(self, model):
+        self.nbtTypesMenu = QtGui.QMenu()
-        self.model = model
+        self.nbtTypesMenu.addAction(NBTIcon(1), self.tr("Byte"), self.addByte)
-        self.proxyModel = proxyModel = NBTFilterProxyModel(self)
+        self.nbtTypesMenu.addAction(NBTIcon(2), self.tr("Short"), self.addShort)
-        proxyModel.setSourceModel(model)
+        self.nbtTypesMenu.addAction(NBTIcon(3), self.tr("Int"), self.addInt)
-        proxyModel.setDynamicSortFilter(True)
+        self.nbtTypesMenu.addAction(NBTIcon(4), self.tr("Long"), self.addLong)
+        self.nbtTypesMenu.addAction(NBTIcon(5), self.tr("Float"), self.addFloat)
+        self.nbtTypesMenu.addAction(NBTIcon(6), self.tr("Double"), self.addDouble)
+        self.nbtTypesMenu.addAction(NBTIcon(8), self.tr("String"), self.addString)
+        self.nbtTypesMenu.addAction(NBTIcon(9), self.tr("List"), self.addList)
+        self.nbtTypesMenu.addAction(NBTIcon(10), self.tr("Compound"), self.addCompound)
 
-        self.treeView.setModel(proxyModel)
+        self.nbtTypesMenu.addAction(NBTIcon(7), self.tr("Byte Array"), self.addByteArray)
+        self.nbtTypesMenu.addAction(NBTIcon(11), self.tr("Int Array"), self.addIntArray)
+        self.nbtTypesMenu.addAction(NBTIcon(12), self.tr("Short Array"), self.addShortArray)
+
+
+    def setRootTag(self, rootTag, keepExpanded=False):
+        if rootTag is self.rootTag:
+            return
+        self.rootTag = rootTag
+        if rootTag is None:
+            self.treeView.setModel(None)
+            self.model = None
+            return
+
+        self.model = NBTTreeModel(rootTag)
+        expanded = []
+        current = None
+        if keepExpanded and self.proxyModel:
+            current = self.proxyModel.data(self.treeView.currentIndex(), NBTPathRole)
+            def addExpanded(parentIndex):
+                for row in range(self.proxyModel.rowCount(parentIndex)):
+                    index = self.proxyModel.index(row, 0, parentIndex)
+                    if self.treeView.isExpanded(index):
+                        expanded.append(self.proxyModel.data(index, NBTPathRole))
+                        addExpanded(index)
+
+            addExpanded(QtCore.QModelIndex())
+
+
+        self.model.dataChanged.connect(self.dataDidChange)
+        self.model.rowsInserted.connect(self.rowsDidInsert)
+        self.model.rowsRemoved.connect(self.rowsDidRemove)
+
+        self.proxyModel = NBTFilterProxyModel(self)
+        self.proxyModel.setSourceModel(self.model)
+        # self.proxyModel.setDynamicSortFilter(True)
+
+        self.treeView.setModel(self.model)
         header = self.treeView.header()
         header.setStretchLastSection(False)
         header.setResizeMode(1, header.ResizeMode.Stretch)
         header.setResizeMode(2, header.ResizeMode.Fixed)
         header.setResizeMode(3, header.ResizeMode.Fixed)
 
+        if keepExpanded:
+            for path in expanded:
+                matches = self.proxyModel.match(self.proxyModel.index(0, 0, QtCore.QModelIndex()),
+                                                NBTPathRole, path, flags=Qt.MatchExactly | Qt.MatchRecursive)
+                for i in matches:
+                    self.treeView.setExpanded(i, True)
+            if current is not None:
+                matches = self.proxyModel.match(self.proxyModel.index(0, 0, QtCore.QModelIndex()),
+                                                NBTPathRole, current, flags=Qt.MatchExactly | Qt.MatchRecursive)
+                if len(matches):
+                    self.treeView.setCurrentIndex(matches[0])
+        else:
+            self.treeView.expandToDepth(0)
         self.treeView.sortByColumn(0, Qt.AscendingOrder)
         self.treeView.resizeColumnToContents(0)
         self.treeView.resizeColumnToContents(1)
         self.treeView.resizeColumnToContents(2)
         self.treeView.resizeColumnToContents(3)
 
-    def itemExpanded(self, index):
+    def itemExpanded(self):
         self.treeView.resizeColumnToContents(0)
 
+    indexAddingTo = None
 
     def itemClicked(self, index):
-        index = self.proxyModel.mapToSource(index)
+        #index = self.proxyModel.mapToSource(index)
         item = self.model.getItem(index)
         if index.column() == 2:
-            if item.isList:
+            if item.isList and item.tag.list_type:
-                self.model.insertRow(item.childCount(), index)
+                row = item.childCount()
-            elif item.isCompound:
+                self.model.insertRow(row, index)
-                """ show tag type menu """
+                newItemIndex = self.model.index(row, 1, index)
+                #self.treeView.setCurrentIndex(self.proxyModel.mapFromSource(newItemIndex))
+                #self.treeView.edit(self.proxyModel.mapFromSource(newItemIndex))
+
+            if item.isCompound or (item.isList and not item.tag.list_type):
+                self.indexAddingTo = index
+                self.nbtTypesMenu.move(QtGui.QCursor.pos())
+                self.nbtTypesMenu.show()
         if index.column() == 3:
             parent = self.model.parent(index)
+            self.doomedTagName = self.tagNameForUndo(index)
             self.model.removeRow(index.row(), parent)
+
+    def addItemWithType(self, tagID):
+        if not self.indexAddingTo:
+            return
+        item = self.model.getItem(self.indexAddingTo)
+        row = item.childCount()
+        self.model.insertRow(row, self.indexAddingTo, tagID)
+        newItemIndex = self.model.index(row, 0 if item.isCompound else 1, self.indexAddingTo)
+        #self.treeView.setCurrentIndex(self.proxyModel.mapFromSource(newItemIndex))
+        #self.treeView.edit(self.proxyModel.mapFromSource(newItemIndex))
+        self.indexAddingTo = None
+
+    def addByte(self):
+        self.addItemWithType(1)
+
+    def addShort(self):
+        self.addItemWithType(2)
+
+    def addInt(self):
+        self.addItemWithType(3)
+
+    def addLong(self):
+        self.addItemWithType(4)
+
+    def addFloat(self):
+        self.addItemWithType(5)
+
+    def addDouble(self):
+        self.addItemWithType(6)
+
+    def addByteArray(self):
+        self.addItemWithType(7)
+
+    def addString(self):
+        self.addItemWithType(8)
+
+    def addList(self):
+        self.addItemWithType(9)
+
+    def addCompound(self):
+        self.addItemWithType(10)
+
+    def addIntArray(self):
+        self.addItemWithType(11)
+
+    def addShortArray(self):
+        self.addItemWithType(12)
+
+
+
+    def tagNameForUndo(self, index):
+        parent = self.model.parent(index)
+        item = self.model.getItem(index)
+        parentItem = self.model.getItem(parent)
+        if parentItem is not None and parentItem.isList:
+            name = "%s #%d" % (self.tagNameForUndo(parent), parentItem.tag.index(item.tag))
+        else:
+            name = item.tag.name
+        return name
+
+    def dataDidChange(self, index):
+        name = self.tagNameForUndo(index)
+        if index.column() == 0:
+            text = "%sRename NBT tag %s" % (self.undoCommandPrefixText, name)
+        elif index.column() == 1:
+            text = "%sChange value of NBT tag %s" % (self.undoCommandPrefixText, name)
+        else:
+            text = "Unknown data changed."
+
+        command = NBTDataChangeCommand(self.editorSession, text)
+        with command.begin():
+            self.editMade.emit()
+            self.editorSession.worldEditor.syncToDisk()
+        self.editorSession.pushCommand(command)
+
+    def rowsDidInsert(self, index):
+        name = self.tagNameForUndo(index.parent())
+        text = "%sInsert NBT tag under %s" % (self.undoCommandPrefixText, name)
+
+        command = NBTDataChangeCommand(self.editorSession, text)
+        with command.begin():
+            self.editMade.emit()
+            self.editorSession.worldEditor.syncToDisk()
+        self.editorSession.pushCommand(command)
+
+    doomedTagName = None
+
+    def rowsDidRemove(self, index, start, end):
+        name = self.tagNameForUndo(index)
+        text = "%sRemove NBT tag %s from %s" % (self.undoCommandPrefixText, self.doomedTagName, name)
+
+        command = NBTDataChangeCommand(self.editorSession, text)
+        with command.begin():
+            self.editMade.emit()
+            self.editorSession.worldEditor.syncToDisk()
+        self.editorSession.pushCommand(command)