Add index functions to extract headers and word counts for a given handle

nw/core/index.py

@@ -608,6 +608,36 @@ class NWIndex():
 
         return hCount
 
+    def getHandleWordCounts(self, tHandle):
+        """Get all header word counts for a specific handle.
+        """
+        theCounts = []
+        hRecord = self._novelIndex.get(tHandle, None)
+        if hRecord is None:
+            hRecord = self._noteIndex.get(tHandle, None)
+        if hRecord is None:
+            return theCounts
+
+        for sTitle, sData in hRecord.items():
+            theCounts.append(("%s:%s" % (tHandle, sTitle), sData["wCount"]))
+
+        return theCounts
+
+    def getHandleHeaders(self, tHandle):
+        """Get all headers for a specific handle.
+        """
+        theHeaders = []
+        hRecord = self._novelIndex.get(tHandle, None)
+        if hRecord is None:
+            hRecord = self._noteIndex.get(tHandle, None)
+        if hRecord is None:
+            return theHeaders
+
+        for sTitle, sData in hRecord.items():
+            theHeaders.append((sTitle, sData["level"], sData["title"]))
+
+        return theHeaders
+
     def getTableOfContents(self, maxDepth, skipExcluded=True):
         """Generate a table of contents up to a maxiumum depth.
         """

tests/test_core/test_core_index.py

@@ -572,12 +572,12 @@ def testCoreIndex_ExtractData(nwMinimal, dummyGUI):
         "@char: Jane\n\n"
         "% this is a comment\n\n"
         "This is a story about Jane Smith.\n\n"
-        "Well, not really.\n"
+        "Well, not really. She's still awesome though.\n"
     ))
     # Whole document
     cC, wC, pC = theIndex.getCounts(nHandle)
-    assert cC == 124
+    assert cC == 152
-    assert wC == 24
+    assert wC == 28
     assert pC == 4
 
     # First part
@@ -588,8 +588,8 @@ def testCoreIndex_ExtractData(nwMinimal, dummyGUI):
 
     # First part
     cC, wC, pC = theIndex.getCounts(nHandle, "T000011")
-    assert cC == 62
+    assert cC == 90
-    assert wC == 12
+    assert wC == 16
     assert pC == 2
 
     # Get section counts for a note file
@@ -605,12 +605,12 @@ def testCoreIndex_ExtractData(nwMinimal, dummyGUI):
         "@char: Jane\n\n"
         "% this is a comment\n\n"
         "This is a story about Jane Smith.\n\n"
-        "Well, not really.\n"
+        "Well, not really. She's still awesome though.\n"
     ))
     # Whole document
     cC, wC, pC = theIndex.getCounts(cHandle)
-    assert cC == 124
+    assert cC == 152
-    assert wC == 24
+    assert wC == 28
     assert pC == 4
 
     # First part
@@ -621,8 +621,8 @@ def testCoreIndex_ExtractData(nwMinimal, dummyGUI):
 
     # First part
     cC, wC, pC = theIndex.getCounts(cHandle, "T000011")
-    assert cC == 62
+    assert cC == 90
-    assert wC == 12
+    assert wC == 16
     assert pC == 2
 
     ##
@@ -650,7 +650,7 @@ def testCoreIndex_ExtractData(nwMinimal, dummyGUI):
     theProject.projTree._treeOrder.remove("0000000000000")
 
     # Extract stats
-    assert theIndex.getNovelWordCount(False) == 30
+    assert theIndex.getNovelWordCount(False) == 34
     assert theIndex.getNovelWordCount(True) == 6
     assert theIndex.getNovelTitleCounts(False) == [0, 2, 1, 2, 0]
     assert theIndex.getNovelTitleCounts(True) == [0, 0, 1, 2, 0]
@@ -670,9 +670,29 @@ def testCoreIndex_ExtractData(nwMinimal, dummyGUI):
     assert theIndex.getTableOfContents(0, False) == []
     assert theIndex.getTableOfContents(1, False) == [
         ("%s:T000001" % nHandle, 1, "Hello World!", 12),
-        ("%s:T000011" % nHandle, 1, "Hello World!", 18),
+        ("%s:T000011" % nHandle, 1, "Hello World!", 22),
+    ]
+
+    # Header Word Counts
+    bHandle = "0000000000000"
+    assert theIndex.getHandleWordCounts(bHandle) == []
+    assert theIndex.getHandleWordCounts(hHandle) == [("%s:T000001" % hHandle, 2)]
+    assert theIndex.getHandleWordCounts(sHandle) == [("%s:T000001" % sHandle, 2)]
+    assert theIndex.getHandleWordCounts(tHandle) == [("%s:T000001" % tHandle, 2)]
+    assert theIndex.getHandleWordCounts(nHandle) == [
+        ("%s:T000001" % nHandle, 12), ("%s:T000011" % nHandle, 16)
     ]
 
     assert theProject.closeProject()
 
+    # Header Record
+    bHandle = "0000000000000"
+    assert theIndex.getHandleHeaders(bHandle) == []
+    assert theIndex.getHandleHeaders(hHandle) == [("T000001", "H2", "Chapter One")]
+    assert theIndex.getHandleHeaders(sHandle) == [("T000001", "H3", "Scene One")]
+    assert theIndex.getHandleHeaders(tHandle) == [("T000001", "H3", "Scene Two")]
+    assert theIndex.getHandleHeaders(nHandle) == [
+        ("T000001", "H1", "Hello World!"), ("T000011", "H1", "Hello World!")
+    ]
+
 # END Test testCoreIndex_ExtractData