fix: fix length check and image order

feat: improve efficiency of restore_toc_list

Author: mddanish00

src/stas_ln_translator/process.py

@@ -30,11 +30,14 @@ def preprocess_document(soup: BeautifulSoup) -> BeautifulSoup:
     # Find all <p> tags
     for p_tag in soup.find_all("p"):
         # Remove <img> wrapped with <p> and add after <p>
-        if len(list(p_tag.children)) == 1 and list(p_tag.children)[0].name == "img":
+        # More robustly check if the paragraph only contains an image, allowing for whitespace.
+        child_tags = p_tag.find_all(True, recursive=False)
+        if not p_tag.get_text(strip=True) and len(child_tags) == 1 and child_tags[0].name == 'img':
             p_tag.unwrap()
             continue
-
-        for img in p_tag.find_all("img"):
+        
+        # Iterate in reverse to maintain the original order of multiple images when moving them after the paragraph.
+        for img in reversed(p_tag.find_all("img")):
             p_tag.insert_after(img.extract())
 
         # Remove <span> tags but preserve their content

src/stas_ln_translator/utils.py

@@ -160,53 +160,28 @@ def restore_toc_list(
     Returns:
         TOCList: The restored table of contents (TOC) list.
     """
-    restored_list: TOCList = []
-    top_list = [i for i in zip(flattened_list, index_list) if i[1].count("-") == 0]
+    assert len(flattened_list) == len(index_list)
 
-    for item, index in top_list:
-        if isinstance(item, epub.Link):
-            restored_list.append(item)
-        elif isinstance(item, epub.Section):
-            restored_list.append(
-                (
-                    item,
-                    _get_children_toc_list(
-                        index,
-                        [
-                            i
-                            for i in zip(flattened_list, index_list)
-                            if i[1].startswith(index)
-                        ],
-                    ),
-                )
-            )
-    return restored_list
+    items_by_index = dict(zip(index_list, flattened_list))
+    children_by_parent_index: dict[str, list[str]] = {}
 
+    for index in index_list:
+        parent_index = "root"
+        if "-" in index:
+            parent_index = index.rpartition("-")[0]
 
-def _get_children_toc_list(
-    parent_index: str,
-    current_list: list[tuple[epub.Link | epub.Section, str]],
-) -> TOCList:
-    """The recursive function to recursively get children TOC list.
-
-    Args:
-        parent_index (str): The parent index of the current list.
-        current_list (list[tuple[epub.Link  |  epub.Section, str]]): The current list of items to process.
+        children_by_parent_index.setdefault(parent_index, []).append(index)
 
-    Returns:
-        TOCList: The restored table of contents (TOC) list.
-    """
-    children: TOCList = []
+    def build_tree(parent_index: str) -> TOCList:
+        restored_list: TOCList = []
+        child_indices = children_by_parent_index.get(parent_index, [])
 
-    for i, (item, index) in enumerate(current_list):
-        if index.startswith(parent_index) and (
-            index.count("-") == parent_index.count("-") + 1
-        ):
+        for child_index in child_indices:
+            item = items_by_index[child_index]
             if isinstance(item, epub.Link):
-                children.append(item)
+                restored_list.append(item)
             elif isinstance(item, epub.Section):
-                # Recursively find children for this section
-                section_children = _get_children_toc_list(index, current_list[i + 1 :])
-                children.append((item, section_children))
+                restored_list.append((item, build_tree(child_index)))
+        return restored_list
 
-    return children
+    return build_tree("root")