From 4c27d2600aa1598e32b4fc77d8423d959c0df80d Mon Sep 17 00:00:00 2001
From: "Veronica K. B. Olsen" <1619840+vkbo@users.noreply.github.com>
Date: Fri, 29 Jan 2021 01:09:04 +0100
Subject: [PATCH] Update document layout during editor save process

---
 nw/core/tree.py     | 48 ++++++++++++++++++++++++++++++++++++++++++++-
 nw/gui/doceditor.py |  8 +++++++-
 2 files changed, 54 insertions(+), 2 deletions(-)

diff --git a/nw/core/tree.py b/nw/core/tree.py
index 36a56e5c..786ba89f 100644
--- a/nw/core/tree.py
+++ b/nw/core/tree.py
@@ -33,10 +33,36 @@ from time import time
 
 from nw.core.item import NWItem
 from nw.common import checkHandle
-from nw.constants import nwFiles, nwItemType, nwItemClass, nwItemLayout, nwConst
+from nw.constants import (
+    nwFiles, nwItemType, nwItemClass, nwItemLayout, nwConst, nwLists
+)
 
 logger = logging.getLogger(__name__)
 
+# Translation map for item layouts
+NOVEL_MAP = {
+    nwItemLayout.SCENE: {
+        "H1": nwItemLayout.PARTITION,
+        "H2": nwItemLayout.CHAPTER,
+        "H4": nwItemLayout.SCENE,
+    },
+    nwItemLayout.CHAPTER: {
+        "H1": nwItemLayout.PARTITION,
+        "H3": nwItemLayout.SCENE,
+        "H4": nwItemLayout.SCENE,
+    },
+    nwItemLayout.UNNUMBERED: {
+        "H1": nwItemLayout.PARTITION,
+        "H3": nwItemLayout.SCENE,
+        "H4": nwItemLayout.SCENE,
+    },
+    nwItemLayout.PARTITION: {
+        "H2": nwItemLayout.CHAPTER,
+        "H3": nwItemLayout.SCENE,
+        "H4": nwItemLayout.SCENE,
+    },
+}
+
 class NWTree():
 
     def __init__(self, theProject):
@@ -202,6 +228,26 @@ class NWTree():
                 novelWords += tItem.wordCount
         return novelWords, noteWords
 
+    def updateItemLayout(self, tHandle, hLevel):
+        """Check if the item layout needs updating based on the header
+        given level.
+        """
+        tItem = self.__getitem__(tHandle)
+        if tItem is None:
+            return False
+        if tItem.itemClass not in nwLists.CLS_NOVEL:
+            return False
+        if hLevel not in ("H1", "H2", "H3", "H4"):
+            return False
+
+        iLayout = tItem.itemLayout
+        if iLayout in NOVEL_MAP:
+            if hLevel in NOVEL_MAP[iLayout]:
+                tItem.itemLayout = NOVEL_MAP[iLayout][hLevel]
+                return True
+
+        return False
+
     ##
     #  Tree Structure Methods
     ##
diff --git a/nw/gui/doceditor.py b/nw/gui/doceditor.py
index b3709420..9990988a 100644
--- a/nw/gui/doceditor.py
+++ b/nw/gui/doceditor.py
@@ -393,6 +393,7 @@ class GuiDocEditor(QTextEdit):
             return False
 
         docText = self.getText()
+        tHandle = theItem.itemHandle
 
         cC, wC, pC = countWords(docText)
         self._updateCounts(cC, wC, pC)
@@ -405,7 +406,12 @@ class GuiDocEditor(QTextEdit):
         self.nwDocument.saveDocument(docText)
         self.setDocumentChanged(False)
 
-        self.theParent.theIndex.scanText(theItem.itemHandle, docText)
+        self.theParent.theIndex.scanText(tHandle, docText)
+
+        hLevel, _ = self.theParent.theIndex.getFirstTitle(tHandle)
+        if self.theProject.projTree.updateItemLayout(tHandle, hLevel):
+            self.theParent.treeView.setTreeItemValues(tHandle)
+            self.nwDocument.saveDocument(docText)
 
         return True