# -*- coding: utf-8 -*- """novelWriter Project Tree Class novelWriter – Project Tree Class ================================== Class holding the data of the project tree File History: Created: 2020-05-07 [0.4.5] This file is a part of novelWriter Copyright 2020, Veronica Berglyd Olsen This program is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details. You should have received a copy of the GNU General Public License along with this program. If not, see . """ import logging import json from os import path from lxml import etree from hashlib import sha256 from time import time from nw.core.item import NWItem from nw.common import checkString from nw.constants import nwFiles, nwItemType, nwItemClass, nwItemLayout, nwConst logger = logging.getLogger(__name__) class NWTree(): def __init__(self, theProject): self.theProject = theProject 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._trashRoot = None # The handle of the trash root folder self._theLength = 0 # Always the length of _treeOrder self._theIndex = 0 # The current iterator index self._treeChanged = False # True if tree structure has changed self._handleSeed = None # Used for generating handles for testing return ## # Class Methods ## def clear(self): """Clear the item tree entirely. """ self._projTree = {} self._treeOrder = [] self._treeRoots = [] self._trashRoot = None self._archRoot = None self._theLength = 0 self._theIndex = 0 self._treeChanged = False return def handles(self): """Returns a copy of the list of all the active handles. """ return self._treeOrder.copy() def append(self, tHandle, pHandle, nwItem): """Add a new item to the end of the tree. """ tHandle = checkString(tHandle, None, True) pHandle = checkString(pHandle, None, True) if tHandle is None: tHandle = self._makeHandle() logger.verbose("Adding entry %s with parent %s" % (str(tHandle), str(pHandle))) nwItem.setHandle(tHandle) nwItem.setParent(pHandle) self._projTree[tHandle] = nwItem self._treeOrder.append(tHandle) if nwItem.itemType == nwItemType.ROOT: logger.verbose("Entry %s is a root item" % str(tHandle)) self._treeRoots.append(tHandle) if nwItem.itemClass == nwItemClass.ARCHIVE: logger.verbose("Entry %s is the archive folder" % str(tHandle)) self._archRoot = tHandle if nwItem.itemType == nwItemType.TRASH: if self._trashRoot is None: logger.verbose("Entry %s is the trash folder" % str(tHandle)) self._trashRoot = tHandle else: logger.error("Only one trash folder allowed") self._theLength = len(self._treeOrder) self._setTreeChanged(True) return def packXML(self, xParent): """Pack the content of the tree into an XML object. """ xContent = etree.SubElement(xParent, "content", attrib={ "count": str(self._theLength)} ) for tHandle in self._treeOrder: tItem = self.__getitem__(tHandle) tItem.packXML(xContent) return def unpackXML(self, xContent): """Iterate through all items of a content XML object and add them to the project tree. """ if xContent.tag != "content": logger.error("XML entry is not a NWTree") return False self.clear() for xItem in xContent: nwItem = NWItem(self.theProject) if nwItem.unpackXML(xItem): self.append(nwItem.itemHandle, nwItem.parHandle, nwItem) nwItem.saveInitialCount() return True def writeToCFiles(self): """Write the convenience table of contents files in the root of the project directory. These files are there to assist the user if they wish to browse the stored files. """ tocText = path.join(self.theProject.projPath, nwFiles.TOC_TXT) tocJson = path.join(self.theProject.projPath, nwFiles.TOC_JSON) jsonData = [] try: # Dump the text with open(tocText, mode="w", encoding="utf8") as outFile: outFile.write("\n") outFile.write(" Table of Contents\n") outFile.write("===================\n") outFile.write("\n") outFile.write(" %-25s %-9s %s\n" % ("File Name", "Class", "Document Label")) outFile.write("-"*80+"\n") for tHandle in sorted(self._treeOrder): tItem = self.__getitem__(tHandle) if tItem is None: continue tFile = tHandle+".nwd" if path.isfile(path.join(self.theProject.projContent, tFile)): outFile.write(" %-25s %-9s %s\n" % ( path.join("content", tFile), tItem.itemClass.name, tItem.itemName, )) jsonData.append([ path.join("content", tFile), tItem.itemClass.name, tItem.itemName, ]) outFile.write("\n") # Dump the JSON with open(tocJson, mode="w+", encoding="utf8") as outFile: outFile.write(json.dumps(jsonData, indent=2)) except Exception as e: logger.error(str(e)) return def sumWords(self): """Loops over all entries and adds up the word counts. """ noteWords = 0 novelWords = 0 for tHandle in self._treeOrder: tItem = self.__getitem__(tHandle) if tItem is None: continue if tItem.itemLayout == nwItemLayout.NO_LAYOUT: pass elif tItem.itemLayout == nwItemLayout.NOTE: noteWords += tItem.wordCount else: novelWords += tItem.wordCount return novelWords, noteWords ## # Tree Structure Methods ## def trashRoot(self): """Returns the handle of the trash folder, or None if there isn't one. """ if self._trashRoot: return self._trashRoot return None def archiveRoot(self): """Returns the handle of the archive folder, or None if there isn't one. """ if self._archRoot: return self._archRoot return None def findRoot(self, theClass): """Find the root item for a given class. Note: This returns the first item for class CUSTOM. """ for aRoot in self._treeRoots: tItem = self.__getitem__(aRoot) if tItem is None: continue if theClass == tItem.itemClass: return tItem.itemHandle return None def checkRootUnique(self, theClass): """Checks if there already is a root entry of class 'theClass' in the root of the project tree. CUSTOM class is skipped as it is not required to be unique. """ if theClass == nwItemClass.CUSTOM: return True for aRoot in self._treeRoots: tItem = self.__getitem__(aRoot) if theClass == tItem.itemClass: return False return True 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. """ tItem = self.__getitem__(tHandle) if tItem is not None: for i in range(nwConst.maxDepth + 1): if tItem.parHandle is None: return tItem else: tHandle = tItem.parHandle tItem = self.__getitem__(tHandle) return None def getItemPath(self, tHandle): """Iterate upwards in the tree until we find the item with parent None, the root item, and return the list of handles. We do this with a for loop with a maximum depth to make infinite loops impossible. """ tTree = [] tItem = self.__getitem__(tHandle) if tItem is not None: tTree.append(tHandle) for i in range(nwConst.maxDepth + 1): if tItem.parHandle is None: return tTree else: tHandle = tItem.parHandle tItem = self.__getitem__(tHandle) if tItem is None: return tTree else: tTree.append(tHandle) return tTree ## # Setters ## def setOrder(self, newOrder): """Reorders the tree based on a list of items. """ tmpOrder = [] # Add all known elements to a new temp list for tHandle in newOrder: if tHandle in self._projTree: tmpOrder.append(tHandle) else: logger.error("Handle %s in new tree order is not in project tree" % tHandle) # Do a reverse lookup to check for items that will be lost # This is mainly for debugging purposes for tHandle in self._treeOrder: if tHandle not in tmpOrder: logger.warning("Handle %s in old tree order is not in new tree order" % tHandle) # Save the temp list self._treeOrder = tmpOrder self._theLength = len(self._treeOrder) self._setTreeChanged(True) logger.verbose("Project tree order updated") return def setSeed(self, theSeed): """Used for debugging! Sets a seed for generating handles so that they always come out in a predictable order. """ self._handleSeed = theSeed return def setFileItemLayout(self, tHandle, itemLayout): """Set the nwItemLayout for a specific file. """ tItem = self.__getitem__(tHandle) if tItem is None: return False if tItem.itemType != nwItemType.FILE: logger.error("Item '%s' is not a file" % tHandle) return False if not isinstance(itemLayout, nwItemLayout): return False tItem.setLayout(itemLayout) return True ## # Getters ## def countTypes(self): """Count the number of files, folders and roots in the project. """ nRoot = 0 nFolder = 0 nFile = 0 for tHandle in self._treeOrder: tItem = self.__getitem__(tHandle) if tItem is None: continue elif tItem.itemType == nwItemType.ROOT: nRoot += 1 elif tItem.itemType == nwItemType.FOLDER: nFolder += 1 elif tItem.itemType == nwItemType.FILE: nFile += 1 return nRoot, nFolder, nFile ## # Meta Methods ## def __len__(self): """Return the length counter. Does not check that it is correct! """ return self._theLength def __bool__(self): """Returns True if the tree has any entries. """ return self._theLength > 0 ## # Item Access Methods ## def __getitem__(self, tHandle): """Return a project item based on its handle. Returns None if the handle doesn't exist in the project. """ if tHandle in self._projTree: return self._projTree[tHandle] logger.error("No tree item with handle %s" % str(tHandle)) return None def __delitem__(self, tHandle): """This only removes the item from the order list, but not from the project tree. """ if tHandle not in self._treeOrder: logger.warning( "Could not remove item %s from project tree as it does not exist" % tHandle ) return False self._treeOrder.remove(tHandle) self._theLength = len(self._treeOrder) if tHandle in self._treeRoots: self._treeRoots.remove(tHandle) self._setTreeChanged(True) return True def __contains__(self, tHandle): """Checks if a handle exists in the tree. """ return tHandle in self._treeOrder ## # Iterator Methods ## def __iter__(self): """Initiates the iterator. """ self._theIndex = 0 return self def __next__(self): """Returns the item from the next entry in the _treeOrder list. """ if self._theIndex < self._theLength: theItem = self.__getitem__(self._treeOrder[self._theIndex]) self._theIndex += 1 return theItem else: raise StopIteration ## # Internal Functions ## def _setTreeChanged(self, theState): """Set the changed flag to theState, and if being set to True, propagate that state change to the parent NWProject class. """ self._treeChanged = theState if theState: self.theProject.setProjectChanged(True) return def _makeHandle(self, addSeed=""): """Generate a unique item handle. In the unlikely event that the key already exists, salt the seed and generate a new handle. """ if self._handleSeed is None: newSeed = str(time()) + addSeed else: # This is used for debugging newSeed = str(self._handleSeed) self._handleSeed += 1 logger.verbose("Generating handle with seed '%s'" % newSeed) itemHandle = sha256(newSeed.encode()).hexdigest()[0:13] if itemHandle in self._projTree: logger.warning("Duplicate handle encountered! Retrying ...") itemHandle = self._makeHandle(addSeed+"!") return itemHandle # END Class NWTree