Restructure the record of root items in the tree class

novelwriter/core/tree.py

@@ -47,7 +47,7 @@ class NWTree():
 
         self._projTree    = {}     # Holds all the items of the project
         self._treeOrder   = []     # The order of the tree items on the tree view
-        self._treeRoots   = []     # The root items of the tree
+        self._treeRoots   = {}     # The root items of the tree
         self._trashRoot   = None   # The handle of the trash root folder
         self._archRoot    = None   # The handle of the archive root folder
         self._theIndex    = 0      # The current iterator index
@@ -67,7 +67,7 @@ class NWTree():
         """
         self._projTree  = {}
         self._treeOrder = []
-        self._treeRoots = []
+        self._treeRoots = {}
         self._trashRoot = None
         self._archRoot  = None
         self._theIndex  = 0
@@ -98,7 +98,7 @@ class NWTree():
 
         if nwItem.itemType == nwItemType.ROOT:
             logger.verbose("Item '%s' is a root item", str(tHandle))
-            self._treeRoots.append(tHandle)
+            self._treeRoots[tHandle] = nwItem
             if nwItem.itemClass == nwItemClass.ARCHIVE:
                 logger.verbose("Item '%s' is the archive folder", str(tHandle))
                 self._archRoot = tHandle
@@ -253,20 +253,34 @@ class NWTree():
                 return tItem.itemHandle
         return None
 
-    def getRootItem(self, tHandle):
-        """Iterate upwards in the tree until we find the item with
-        parent None, the root item. We do this with a for loop with a
-        maximum depth to make infinite loops impossible.
+    def isRoot(self, tHandle):
+        """Check if a handle is a root item.
+        """
+        return tHandle in self._treeRoots
+
+    def updateItemRoot(self, tHandle):
+        """Update the root item handle of a given item.
+        """
+        tItem = self.__getitem__(tHandle)
+        iItem = tItem
+        if iItem is not None:
+            for _ in range(nwConst.MAX_DEPTH + 1):
+                if iItem.itemParent is None:
+                    tItem.setRoot(iItem.itemHandle)
+                    return iItem.itemHandle
+                else:
+                    tHandle = iItem.itemParent
+                    iItem = self.__getitem__(tHandle)
+        return None
+
+    def getItemClass(self, tHandle):
+        """Return the class of a given item.
         """
         tItem = self.__getitem__(tHandle)
         if tItem is not None:
-            for i in range(nwConst.MAX_DEPTH + 1):
-                if tItem.itemParent is None:
-                    return tItem
-                else:
-                    tHandle = tItem.itemParent
-                    tItem = self.__getitem__(tHandle)
-        return None
+            if tItem.itemRoot in self._treeRoots:
+                return self._treeRoots[tItem.itemRoot].itemClass
+        return nwItemClass.NO_CLASS
 
     def getItemPath(self, tHandle):
         """Iterate upwards in the tree until we find the item with
@@ -405,7 +419,7 @@ class NWTree():
             return
 
         if tHandle in self._treeRoots:
-            self._treeRoots.remove(tHandle)
+            del self._treeRoots[tHandle]
         if tHandle == self._trashRoot:
             self._trashRoot = None
         if tHandle == self._archRoot: