Veronica K. B. Olsen
468 lines
15 KiB
Python
468 lines
15 KiB
Python
# -*- 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 <https://www.gnu.org/licenses/>.
|
||
"""
|
||
|
||
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
|
||
|
||
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 of 200 to make infinite loops impossible.
|
||
"""
|
||
tItem = self.__getitem__(tHandle)
|
||
if tItem is not None:
|
||
for i in range(200):
|
||
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 of 200 to make
|
||
infinite loops impossible.
|
||
"""
|
||
tTree = []
|
||
tItem = self.__getitem__(tHandle)
|
||
if tItem is not None:
|
||
tTree.append(tHandle)
|
||
for i in range(200):
|
||
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
|