From af21f3c5d8e86f7087b214a3e548ea810f028070 Mon Sep 17 00:00:00 2001
From: "Veronica K. B. Olsen" <1619840+vkbo@users.noreply.github.com>
Date: Mon, 1 Feb 2021 22:48:07 +0100
Subject: [PATCH] Add index functions to extract headers and word counts for a
 given handle

---
 nw/core/index.py                   | 30 ++++++++++++++++++++
 tests/test_core/test_core_index.py | 44 ++++++++++++++++++++++--------
 2 files changed, 62 insertions(+), 12 deletions(-)

diff --git a/nw/core/index.py b/nw/core/index.py
index ba2e06f3..2299ab0d 100644
--- a/nw/core/index.py
+++ b/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.
         """
diff --git a/tests/test_core/test_core_index.py b/tests/test_core/test_core_index.py
index e10000d6..eeee2be9 100644
--- a/tests/test_core/test_core_index.py
+++ b/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 wC == 24
+    assert cC == 152
+    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 wC == 12
+    assert cC == 90
+    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 wC == 24
+    assert cC == 152
+    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 wC == 12
+    assert cC == 90
+    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