epub-translator 0.1.1__py3-none-any.whl → 0.1.3__py3-none-any.whl

epub_translator/segment/inline_segment.py

@@ -0,0 +1,356 @@
+from collections.abc import Generator, Iterable, Iterator
+from dataclasses import dataclass
+from xml.etree.ElementTree import Element
+
+from ..utils import ensure_list, is_the_same, nest
+from ..xml import ID_KEY, append_text_in_element, iter_with_stack, plain_text
+from .common import FoundInvalidIDError, validate_id_in_element
+from .text_segment import TextSegment
+from .utils import IDGenerator, element_fingerprint, id_in_element
+
+
+@dataclass
+class InlineLostIDError:
+    element: Element
+    stack: list[Element]
+
+
+@dataclass
+class InlineUnexpectedIDError:
+    id: int
+    element: Element
+
+
+@dataclass
+class InlineExpectedIDsError:
+    id2element: dict[int, Element]
+
+
+@dataclass
+class InlineWrongTagCountError:
+    expected_count: int
+    found_elements: list[Element]
+    stack: list[Element]
+
+
+InlineError = InlineLostIDError | InlineUnexpectedIDError | InlineExpectedIDsError | InlineWrongTagCountError
+
+
+def search_inline_segments(text_segments: Iterable[TextSegment]) -> Generator["InlineSegment", None, None]:
+    stack_data: tuple[list[list[TextSegment | InlineSegment]], Element, int] | None = None
+    inline_segment: InlineSegment | None = None
+
+    for text_segment in text_segments:
+        if stack_data is not None:
+            stack, stack_block, stack_base_depth = stack_data
+            if stack_block is not text_segment.block_parent:
+                inline_segment = _pop_stack_data(stack_data)
+                stack_data = None
+                if inline_segment:
+                    yield inline_segment
+
+        if stack_data is None:
+            stack_data = (
+                [],
+                text_segment.block_parent,
+                text_segment.block_depth,
+            )
+
+        stack, stack_block, stack_base_depth = stack_data
+
+        while len(stack) < text_segment.depth + 1:
+            stack.append([])
+
+        while len(stack) > text_segment.depth + 1:
+            _pop_stack(
+                stack=stack,
+                stack_base_depth=stack_base_depth,
+            )
+
+        # text_segment.depth 可视为它在 stack 中的 index，必须令 len(stack) == text_segment.depth + 1
+        stack[-1].append(text_segment)
+
+    if stack_data is not None:
+        inline_segment = _pop_stack_data(stack_data)
+        if inline_segment:
+            yield inline_segment
+
+
+def _pop_stack_data(stack_data: tuple[list[list["TextSegment | InlineSegment"]], Element, int]):
+    stack, _, stack_base_depth = stack_data
+    inline_segment: InlineSegment | None = None
+    while stack:
+        inline_segment = _pop_stack(
+            stack=stack,
+            stack_base_depth=stack_base_depth,
+        )
+    return inline_segment
+
+
+def _pop_stack(
+    stack: list[list["TextSegment | InlineSegment"]],
+    stack_base_depth: int,
+) -> "InlineSegment | None":
+    inline_segment: InlineSegment | None = None
+    depth = len(stack) + stack_base_depth - 1
+    popped = stack.pop()
+    if popped:
+        inline_segment = InlineSegment(depth, popped)
+    if stack and inline_segment is not None:
+        stack[-1].append(inline_segment)
+    return inline_segment
+
+
+class InlineSegment:
+    def __init__(self, depth: int, children: list["TextSegment | InlineSegment"]) -> None:
+        assert depth > 0
+        self.id: int | None = None
+        self._children: list[TextSegment | InlineSegment] = children
+        self._parent_stack: list[Element] = children[0].parent_stack[:depth]
+
+        # 每一组 tag 都对应一个 ids 列表。
+        # 若为空，说明该 tag 属性结构全同，没必要分配 id 以区分。
+        # 若非空，则表示 tag 下每一个 element 都有 id 属性。
+        # 注意，相同 tag 下的 element 要么全部有 id，要么全部都没有 id
+        self._child_tag2ids: dict[str, list[int]] = {}
+        self._child_tag2count: dict[str, int] = {}
+
+        next_temp_id: int = 0
+        terms = nest((child.parent.tag, child) for child in children if isinstance(child, InlineSegment))
+
+        for tag, child_terms in terms.items():
+            self._child_tag2count[tag] = len(child_terms)
+            if not is_the_same(  # 仅当 tag 彼此无法区分时才分配 id，以尽可能减少 id 的数量
+                elements=(element_fingerprint(t.parent) for t in child_terms),
+            ):
+                for child in child_terms:
+                    child.id = next_temp_id
+                    next_temp_id += 1
+
+    @property
+    def head(self) -> TextSegment:
+        first_child = self._children[0]
+        if isinstance(first_child, TextSegment):
+            return first_child
+        else:
+            return first_child.head
+
+    @property
+    def tail(self) -> TextSegment:
+        last_child = self._children[-1]
+        if isinstance(last_child, TextSegment):
+            return last_child
+        else:
+            return last_child.tail
+
+    @property
+    def children(self) -> list["TextSegment | InlineSegment"]:
+        return self._children
+
+    @property
+    def parent(self) -> Element:
+        return self._parent_stack[-1]
+
+    @property
+    def parent_stack(self) -> list[Element]:
+        return self._parent_stack
+
+    def __iter__(self) -> Iterator[TextSegment]:
+        for child in self._children:
+            if isinstance(child, TextSegment):
+                yield child
+            elif isinstance(child, InlineSegment):
+                yield from child
+
+    def recreate_ids(self, id_generator: IDGenerator) -> None:
+        self._child_tag2count.clear()
+        self._child_tag2ids.clear()
+
+        for child in self._children:
+            if isinstance(child, InlineSegment):
+                child_tag = child.parent.tag
+                ids = ensure_list(self._child_tag2ids, child_tag)
+                if child.id is not None:
+                    child.id = id_generator.next_id()
+                    ids.append(child.id)
+                child.recreate_ids(id_generator)
+                self._child_tag2count[child_tag] = self._child_tag2count.get(child_tag, 0) + 1
+
+    def create_element(self) -> Element:
+        element = Element(self.parent.tag)
+        previous_element: Element | None = None
+        for child in self._children:
+            if isinstance(child, InlineSegment):
+                previous_element = child.create_element()
+                element.append(previous_element)
+
+            elif isinstance(child, TextSegment):
+                if previous_element is None:
+                    element.text = append_text_in_element(
+                        origin_text=element.text,
+                        append_text=child.text,
+                    )
+                else:
+                    previous_element.tail = append_text_in_element(
+                        origin_text=previous_element.tail,
+                        append_text=child.text,
+                    )
+        if self.id is not None:
+            element.set(ID_KEY, str(self.id))
+        return element
+
+    def validate(self, validated_element: Element) -> Generator[InlineError | FoundInvalidIDError, None, None]:
+        remain_expected_elements: dict[int, Element] = {}
+        for child in self._child_inline_segments():
+            if child.id is not None:
+                remain_expected_elements[child.id] = child.parent
+
+        for _, child_element in iter_with_stack(validated_element):
+            if child_element is validated_element:
+                continue  # skip the root self
+
+            element_id = id_in_element(child_element)
+            if element_id is None:
+                validated_id = validate_id_in_element(
+                    element=child_element,
+                    enable_no_id=True,
+                )
+                if isinstance(validated_id, FoundInvalidIDError):
+                    yield validated_id
+                continue
+
+            remain_expected_element = remain_expected_elements.pop(element_id, None)
+            if remain_expected_element is None:
+                yield InlineUnexpectedIDError(
+                    id=element_id,
+                    element=child_element,
+                )
+
+        if remain_expected_elements:
+            yield InlineExpectedIDsError(
+                id2element=remain_expected_elements,
+            )
+
+        yield from self._validate_children_structure(validated_element)
+
+    def _child_inline_segments(self) -> Generator["InlineSegment", None, None]:
+        for child in self._children:
+            if isinstance(child, InlineSegment):
+                yield child
+                yield from child._child_inline_segments()  # pylint: disable=protected-access
+
+    def _validate_children_structure(self, validated_element: Element):
+        tag2found_elements: dict[str, list[Element]] = {}
+
+        for child_element in validated_element:
+            ids = self._child_tag2ids.get(child_element.tag, None)
+            if not ids:
+                found_elements = ensure_list(tag2found_elements, child_element.tag)
+                found_elements.append(child_element)
+            else:
+                id_str = child_element.get(ID_KEY, None)
+                if id_str is None:
+                    yield InlineLostIDError(
+                        element=child_element,
+                        stack=[self.parent],
+                    )
+
+        for tag, found_elements in tag2found_elements.items():
+            expected_count = self._child_tag2count.get(tag, 0)
+            if len(found_elements) != expected_count:
+                yield InlineWrongTagCountError(
+                    expected_count=expected_count,
+                    found_elements=found_elements,
+                    stack=[self.parent],
+                )
+
+        for child, child_element in self._match_children(validated_element):
+            # pylint: disable=protected-access
+            for error in child._validate_children_structure(child_element):
+                error.stack.insert(0, self.parent)
+                yield error
+
+    # 即便 self.validate(...) 的错误没有排除干净，也要尽可能匹配一个质量较高（尽力而为）的版本
+    def assign_attributes(self, template_element: Element) -> Element:
+        assigned_element = Element(self.parent.tag, self.parent.attrib)
+        if template_element.text and template_element.text.strip():
+            assigned_element.text = append_text_in_element(
+                origin_text=assigned_element.text,
+                append_text=template_element.text,
+            )
+
+        matched_child_element_ids: set[int] = set()
+        for child, child_element in self._match_children(template_element):
+            child_assigned_element = child.assign_attributes(child_element)
+            assigned_element.append(child_assigned_element)
+            matched_child_element_ids.add(id(child_element))
+
+        assigned_child_element_stack = list(assigned_element)
+        assigned_child_element_stack.reverse()
+
+        previous_assigned_child_element: Element | None = None
+        for child_element in template_element:
+            # 只关心 child_element 是否是分割点，不关心它真实对应。极端情况下可能乱序，只好大致对上就行
+            child_text: str = ""
+            if id(child_element) not in matched_child_element_ids:
+                child_text = plain_text(child_element)
+            elif assigned_child_element_stack:
+                previous_assigned_child_element = assigned_child_element_stack.pop()
+            if child_element.tail is not None:
+                child_text += child_element.tail
+            if not child_text.strip():
+                continue
+            if previous_assigned_child_element is None:
+                assigned_element.text = append_text_in_element(
+                    origin_text=assigned_element.text,
+                    append_text=child_text,
+                )
+            else:
+                previous_assigned_child_element.tail = append_text_in_element(
+                    origin_text=previous_assigned_child_element.tail,
+                    append_text=child_text,
+                )
+        return assigned_element
+
+    def _match_children(self, element: Element) -> Generator[tuple["InlineSegment", Element], None, None]:
+        tag2elements = nest((c.tag, c) for c in element)
+        tag2children = nest(
+            (c.parent.tag, (i, c)) for i, c in enumerate(c for c in self._children if isinstance(c, InlineSegment))
+        )
+        used_ids: set[int] = set()
+        children_and_elements: list[tuple[int, InlineSegment, Element]] = []
+
+        for tag, orders_and_children in tag2children.items():
+            # 优先考虑 id 匹配，剩下的以自然顺序尽可能匹配
+            ids = self._child_tag2ids.get(tag, [])
+            matched_children_elements: list[Element | None] = [None] * len(orders_and_children)
+            not_matched_elements: list[Element] = []
+
+            for child_element in tag2elements.get(tag, []):
+                id_order: int | None = None
+                child_id = id_in_element(child_element)
+                if child_id is not None and child_id not in used_ids:
+                    used_ids.add(child_id)  # 一个 id 只能用一次，防止重复
+                    try:
+                        id_order = ids.index(child_id)
+                    except ValueError:
+                        pass
+                if id_order is None:
+                    not_matched_elements.append(child_element)
+                else:
+                    matched_children_elements[id_order] = child_element
+
+            not_matched_elements.reverse()
+            for i in range(len(matched_children_elements)):
+                if not not_matched_elements:
+                    break
+                matched_element = matched_children_elements[i]
+                if matched_element is None:
+                    matched_children_elements[i] = not_matched_elements.pop()
+
+            for (order, child), child_element in zip(orders_and_children, matched_children_elements):
+                if child_element is not None:
+                    children_and_elements.append((order, child, child_element))
+
+        for _, child, child_element in sorted(children_and_elements, key=lambda x: x[0]):
+            yield child, child_element

epub_translator/{xml_translator → segment}/text_segment.py

@@ -4,13 +4,13 @@ from enum import Enum, auto
 from typing import Self
 from xml.etree.ElementTree import Element
 
-from .utils import expand_left_element_texts, expand_right_element_texts, normalize_text_in_element
+from ..xml import expand_left_element_texts, expand_right_element_texts, normalize_text_in_element
 
 # HTML inline-level elements
 # Reference: https://developer.mozilla.org/en-US/docs/Web/HTML/Inline_elements
 # Reference: https://developer.mozilla.org/en-US/docs/Glossary/Inline-level_content
 _HTML_INLINE_TAGS = frozenset(
-    [
+    (
         # Inline text semantics
         "a",
         "abbr",
@@ -67,7 +67,7 @@ _HTML_INLINE_TAGS = frozenset(
         "del",
         "ins",
         "slot",
-    ]
+    )
 )
 
 
@@ -79,7 +79,6 @@ class TextPosition(Enum):
 @dataclass
 class TextSegment:
     text: str
-    index: int  # *.text is 0, the first *.tail is 1, and so on
     parent_stack: list[Element]
     left_common_depth: int
     right_common_depth: int
@@ -90,6 +89,10 @@ class TextSegment:
     def root(self) -> Element:
         return self.parent_stack[0]
 
+    @property
+    def depth(self) -> int:
+        return len(self.parent_stack) - self.block_depth
+
     @property
     def block_parent(self) -> Element:
         return self.parent_stack[self.block_depth - 1]
@@ -106,7 +109,6 @@ class TextSegment:
     def clone(self) -> "TextSegment":
         return TextSegment(
             text=self.text,
-            index=self.index,
             parent_stack=list(self.parent_stack),
             left_common_depth=self.left_common_depth,
             right_common_depth=self.right_common_depth,
@@ -171,20 +173,18 @@ def _search_text_segments(stack: list[Element], element: Element) -> Generator[T
     if text is not None:
         yield TextSegment(
             text=text,
-            index=0,
             parent_stack=next_stack,
             left_common_depth=0,
             right_common_depth=0,
             block_depth=next_block_depth,
             position=TextPosition.TEXT,
         )
-    for i, child_element in enumerate(element):
+    for child_element in element:
         yield from _search_text_segments(next_stack, child_element)
         child_tail = normalize_text_in_element(child_element.tail)
         if child_tail is not None:
             yield TextSegment(
                 text=child_tail,
-                index=i + 1,
                 parent_stack=next_stack,
                 left_common_depth=0,
                 right_common_depth=0,

epub_translator/segment/utils.py

@@ -0,0 +1,43 @@
+from xml.etree.ElementTree import Element
+
+from ..xml import ID_KEY
+
+
+def element_fingerprint(element: Element) -> str:
+    attrs = sorted(f"{key}={value}" for key, value in element.attrib.items())
+    return f"<{element.tag} {' '.join(attrs)}/>"
+
+
+def unwrap_parents(element: Element) -> tuple[Element, list[Element]]:
+    parents: list[Element] = []
+    while True:
+        if len(element) != 1:
+            break
+        child = element[0]
+        if not element.text:
+            break
+        if not child.tail:
+            break
+        parents.append(element)
+        element = child
+        element.tail = None
+    return element, parents
+
+
+def id_in_element(element: Element) -> int | None:
+    id_str = element.get(ID_KEY, None)
+    if id_str is None:
+        return None
+    try:
+        return int(id_str)
+    except ValueError:
+        return None
+
+
+class IDGenerator:
+    def __init__(self):
+        self._previous_id: int = 0
+
+    def next_id(self) -> int:
+        self._previous_id += 1
+        return self._previous_id