Update handler for update block refactor, handle section headers

Author: isaac091

machine/corpora/place_markers_scripture_update_block_handler.py

@@ -1,82 +1,98 @@
 from __future__ import annotations
 
-from copy import copy
 from typing import List, Sequence
 
 from ..jobs.translation_file_service import PretranslationInfo
 from ..tokenization import LatinWordTokenizer
 from ..translation import WordAlignmentMatrix
 from .aligned_word_pair import AlignedWordPair
-from .scripture_update_block import ScriptureUpdateBlock
-from .scripture_update_block_handler import ScriptureUpdateBlockHandler
-from .scripture_update_element import ScriptureUpdateElement, ScriptureUpdateElementType
+from .usfm_stylesheet import UsfmStylesheet
+from .usfm_tag import UsfmTextType
 from .usfm_token import UsfmToken, UsfmTokenType
+from .usfm_update_block import UsfmUpdateBlock
+from .usfm_update_block_element import UsfmUpdateBlockElement, UsfmUpdateBlockElementType
+from .usfm_update_block_handler import UsfmUpdateBlockHandler
 
 TOKENIZER = LatinWordTokenizer()
+STYLESHEET = UsfmStylesheet("usfm.sty")
 
 
-class PlaceMarkersScriptureUpdateBlockHandler(ScriptureUpdateBlockHandler):
+class PlaceMarkersScriptureUpdateBlockHandler(UsfmUpdateBlockHandler):
 
     def __init__(self, pt_info: Sequence[PretranslationInfo]):
         # TODO: when will len(refs) be >1?
         self._pt_info = {info["refs"][0]: info for info in pt_info}
 
-    def process_block(self, block: ScriptureUpdateBlock) -> ScriptureUpdateBlock:
+    def process_block(self, block: UsfmUpdateBlock) -> UsfmUpdateBlock:
         # Nothing to do if there are no markers to place, no alignment to use, or if the block represents an embed
         if (
             len(block.elements) == 0
-            or str(block.ref) not in self._pt_info.keys()
-            or len(self._pt_info[str(block.ref)]["alignment"]) == 0
-            or block.elements[0].type == ScriptureUpdateElementType.EMBED
+            or str(block.refs[0]) not in self._pt_info.keys()
+            or len(self._pt_info[str(block.refs[0])]["alignment"]) == 0
+            or block.elements[0].type == UsfmUpdateBlockElementType.EMBED
             or not any(
                 (
-                    element.type in [ScriptureUpdateElementType.PARAGRAPH, ScriptureUpdateElementType.STYLE]
+                    element.type in [UsfmUpdateBlockElementType.PARAGRAPH, UsfmUpdateBlockElementType.STYLE]
                     and not element.marked_for_removal
                 )
                 for element in block.elements[1:]
             )
         ):
             return block
 
-        # Parsing the block's elements potentially involves removing elements so they are not processed twice,
-        # but the original block may need to be returned if the two versions of the source/target text to not match up
-        orig_elements = copy(block.elements)
+        # Work on a copy in case the block needs to be returned unchanged
+        orig_elements = list(block.elements)
 
         src_sent = ""
         trg_sent = ""
         to_place = []
         src_marker_idxs = []
         placed_elements = [orig_elements[0]]
-        end_elements = []
         ignored_elements = []
 
+        # Section headers should be ignored but re-inserted in the same position relative to other paragraph markers
+        header_elements = []
+        para_markers_left = 0
+        for i, element in reversed(list(enumerate(orig_elements))):
+            if element.type == UsfmUpdateBlockElementType.PARAGRAPH and not element.marked_for_removal:
+                if STYLESHEET.get_tag(str(element.tokens[0].marker)).text_type == UsfmTextType.SECTION:
+                    if i < len(orig_elements) - 1 and orig_elements[i + 1].type == UsfmUpdateBlockElementType.TEXT:
+                        header_elements.insert(0, (para_markers_left, [element, orig_elements.pop(i + 1)]))
+                    else:
+                        header_elements.insert(0, (para_markers_left, [element]))
+                    orig_elements.pop(i)
+                else:
+                    para_markers_left += 1
+
         # Paragraph markers at the end of the block should stay there
+        end_elements = []
         for i, element in reversed(list(enumerate(orig_elements))):
-            if element.type == ScriptureUpdateElementType.PARAGRAPH:
+            if element.type == UsfmUpdateBlockElementType.PARAGRAPH and not element.marked_for_removal:
                 end_elements.insert(0, element)
                 orig_elements.pop(i)
-            elif element.type != ScriptureUpdateElementType.EMBED_BLOCK:
+            elif element.type != UsfmUpdateBlockElementType.EMBED:
                 break
 
         for element in orig_elements[1:]:
-            if element.type == ScriptureUpdateElementType.EXISTING_TEXT:
-                src_sent += element.tokens[0].to_usfm()
-            if element.type == ScriptureUpdateElementType.INSERTED_TEXT:
-                trg_sent += element.tokens[0].to_usfm()
+            if element.type == UsfmUpdateBlockElementType.TEXT:
+                if element.marked_for_removal:
+                    src_sent += element.tokens[0].to_usfm()
+                else:
+                    trg_sent += element.tokens[0].to_usfm()
 
-            if element.marked_for_removal or element.type == ScriptureUpdateElementType.EMBED_BLOCK:
+            if element.marked_for_removal or element.type == UsfmUpdateBlockElementType.EMBED:
                 ignored_elements.append(element)
-            elif element.type in [ScriptureUpdateElementType.PARAGRAPH, ScriptureUpdateElementType.STYLE]:
+            elif element.type in [UsfmUpdateBlockElementType.PARAGRAPH, UsfmUpdateBlockElementType.STYLE]:
                 to_place.append(element)
                 src_marker_idxs.append(len(src_sent))
 
-        src_toks = self._pt_info[str(block.ref)]["source_toks"]
-        trg_toks = self._pt_info[str(block.ref)]["translation_toks"]
+        src_toks = self._pt_info[str(block.refs[0])]["source_toks"]
+        trg_toks = self._pt_info[str(block.refs[0])]["translation_toks"]
 
         # Don't do anything if the source sentence or pretranslation has changed
         if (
             list(t for t in TOKENIZER.tokenize(src_sent)) != src_toks
-            or list(t for t in TOKENIZER.tokenize(trg_sent)) != trg_toks  # could just use translation for trg
+            or list(t for t in TOKENIZER.tokenize(trg_sent)) != trg_toks
         ):
             return block
 
@@ -98,7 +114,7 @@ def process_block(self, block: ScriptureUpdateBlock) -> ScriptureUpdateBlock:
                 if i == 0:
                     adj_src_toks.append(i)
 
-        alignment = to_word_alignment_matrix(self._pt_info[str(block.ref)]["alignment"])
+        alignment = to_word_alignment_matrix(self._pt_info[str(block.refs[0])]["alignment"])
         adj_trg_toks = [
             self._predict_marker_location(alignment, adj_src_tok, src_toks, trg_toks) for adj_src_tok in adj_src_toks
         ]
@@ -115,26 +131,39 @@ def process_block(self, block: ScriptureUpdateBlock) -> ScriptureUpdateBlock:
             to_insert.insert(insert_pos, (trg_str_idx, element))
 
         # Construct new text tokens to put between markers
-        placed_elements.append(
-            ScriptureUpdateElement(
-                ScriptureUpdateElementType.INSERTED_TEXT,
-                [UsfmToken(UsfmTokenType.TEXT, text=trg_sent[: to_insert[0][0]] if len(to_insert) > 0 else trg_sent)],
+        # and reincorporate headers and empty end-of-verse paragraph markers
+        if len(to_insert) == 0 or to_insert[0][0] > 0:
+            placed_elements.append(
+                UsfmUpdateBlockElement(
+                    UsfmUpdateBlockElementType.TEXT,
+                    [
+                        UsfmToken(
+                            UsfmTokenType.TEXT, text=trg_sent[: to_insert[0][0]] if len(to_insert) > 0 else trg_sent
+                        )
+                    ],
+                )
             )
-        )
         for j, (insert_idx, element) in enumerate(to_insert):
+            if element.type == UsfmUpdateBlockElementType.PARAGRAPH:
+                while len(header_elements) > 0 and header_elements[0][0] == para_markers_left:
+                    placed_elements += header_elements.pop(0)[1]
+                para_markers_left -= 1
+
             placed_elements.append(element)
             text_token = UsfmToken(
                 UsfmTokenType.TEXT,
                 text=(trg_sent[insert_idx : to_insert[j + 1][0]] if j + 1 < len(to_insert) else trg_sent[insert_idx:]),
             )
-            placed_elements.append(
-                ScriptureUpdateElement(
-                    ScriptureUpdateElementType.INSERTED_TEXT,
-                    [text_token],
-                )
-            )
+            placed_elements.append(UsfmUpdateBlockElement(UsfmUpdateBlockElementType.TEXT, [text_token]))
+        for element in end_elements:
+            while len(header_elements) > 0 and header_elements[0][0] == para_markers_left:
+                placed_elements += header_elements.pop(0)[1]
+            para_markers_left -= 1
+            placed_elements.append(element)
+        while len(header_elements) > 0:
+            placed_elements += header_elements.pop(0)[1]
 
-        block._elements = placed_elements + end_elements + ignored_elements
+        block._elements = placed_elements + ignored_elements
         return block
 
     def _predict_marker_location(