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

epub_translator/translator.py

@@ -1,184 +1,214 @@
-from os import PathLike
+from collections.abc import Callable
 from pathlib import Path
-from enum import auto, Enum
-from tempfile import mkdtemp
-from shutil import rmtree
+from xml.etree.ElementTree import Element
 
+from .epub import Placeholder, Zip, is_placeholder_tag, read_toc, search_spine_paths, write_toc
+from .epub.common import find_opf_path
 from .llm import LLM
-from .epub import HTMLFile
-from .zip_context import ZipContext
-from .translation import translate as _translate, Incision, Fragment, Language, ProgressReporter
+from .xml import XMLLikeNode, deduplicate_ids_in_element, find_first, plain_text
+from .xml_translator import XMLGroupContext, XMLTranslator, submit_text_segments
 
 
-class TranslatedWriteMode(Enum):
-  APPEND = auto()
-  REPLACE = auto()
-
 def translate(
-      llm: LLM,
-      source_path: PathLike,
-      translated_path: PathLike,
-      target_language: Language,
-      write_mode: TranslatedWriteMode = TranslatedWriteMode.APPEND,
-      user_prompt: str | None = None,
-      working_path: PathLike | None = None,
-      max_chunk_tokens_count: int = 3000,
-      max_threads_count: int = 1,
-      report_progress: ProgressReporter | None = None,
-    ) -> None:
-
-  source_path = Path(source_path)
-  translated_path = Path(translated_path)
-  working_path = Path(working_path) if working_path else None
-  report_progress = report_progress or (lambda _: None)
-
-  _Translator(
-    llm=llm,
-    target_language=target_language,
-    write_mode=write_mode,
-    user_prompt=user_prompt,
-    max_chunk_tokens_count=max_chunk_tokens_count,
-    max_threads_count=max_threads_count,
-    report_progress=report_progress,
-  ).do(
-    source_path=source_path,
-    translated_path=translated_path,
-    working_path=working_path,
-  )
-
-class _Translator:
-  def __init__(
-        self,
-        llm: LLM,
-        target_language: Language,
-        write_mode: TranslatedWriteMode,
-        user_prompt: str | None,
-        max_chunk_tokens_count: int,
-        max_threads_count: int,
-        report_progress: ProgressReporter,
-      ) -> None:
-
-    self._llm: LLM = llm
-    self._target_language: Language = target_language
-    self._write_mode: TranslatedWriteMode = write_mode
-    self._user_prompt: str | None = user_prompt
-    self._max_chunk_tokens_count: int = max_chunk_tokens_count
-    self._max_threads_count: int = max_threads_count
-    self._report_progress: ProgressReporter = report_progress
-
-  def do(self, source_path: Path, translated_path: Path, working_path: Path | None) -> None:
-    is_temp_workspace = not bool(working_path)
-    working_path = working_path or Path(mkdtemp())
-    try:
-      temp_dir = _clean_path(working_path / "temp")
-      temp_dir.mkdir(parents=True, exist_ok=True)
-      cache_path = working_path / "cache"
-
-      context = ZipContext(
-        epub_path=Path(source_path),
-        temp_dir=temp_dir,
-      )
-      context.replace_ncx(lambda texts: self._translate_ncx(
-        texts=texts,
-        cache_path=cache_path,
-        report_progress=lambda p: self._report_progress(p * 0.1)),
-      )
-      self._translate_spine(
-        context=context,
-        cache_path=cache_path,
-        report_progress=lambda p: self._report_progress(0.1 + p * 0.8),
-      )
-      context.archive(translated_path)
-      self._report_progress(1.0)
-
-    finally:
-      if is_temp_workspace:
-        rmtree(working_path, ignore_errors=True)
-
-  def _translate_ncx(self, texts: list[str], cache_path: Path, report_progress: ProgressReporter) -> list[str]:
-    return list(_translate(
-      llm=self._llm,
-      cache_path=cache_path,
-      max_chunk_tokens_count=self._max_chunk_tokens_count,
-      max_threads_count=1,
-      target_language=self._target_language,
-      user_prompt=self._user_prompt,
-      report_progress=report_progress,
-      gen_fragments_iter=lambda: (
-        Fragment(
-          text=text,
-          start_incision=Incision.IMPOSSIBLE,
-          end_incision=Incision.IMPOSSIBLE,
-        )
-        for text in texts
-      ),
-    ))
-
-  def _translate_spine(self, context: ZipContext, cache_path: Path, report_progress: ProgressReporter):
-    spine_paths_iter = iter(list(context.search_spine_paths()))
-    spine: tuple[Path, HTMLFile] | None = None
-    translated_texts: list[str] = []
-    translated_count: int = 0
-    append = (self._write_mode == TranslatedWriteMode.APPEND)
-
-    for translated_text in _translate(
-      llm=self._llm,
-      gen_fragments_iter=lambda: _gen_fragments(context),
-      cache_path=cache_path,
-      max_chunk_tokens_count=self._max_chunk_tokens_count,
-      max_threads_count=self._max_threads_count,
-      target_language=self._target_language,
-      user_prompt=self._user_prompt,
-      report_progress=report_progress,
-    ):
-      did_touch_end = False
-
-      if spine and translated_count >= len(translated_texts):
-        spine_path, spine_file = spine
-        spine_file.write_texts(translated_texts, append)
-        context.write_spine_file(spine_path, spine_file)
-        spine = None
-
-      while not spine:
-        spine_path = next(spine_paths_iter, None)
-        if spine_path is None:
-          spine = None
-          did_touch_end = True
-          break
-        spine_file = context.read_spine_file(spine_path)
-        if spine_file.texts_length == 0:
-          continue
-        spine = (spine_path, spine_file)
-        translated_texts = [""] * spine_file.texts_length
-        translated_count = 0
-        break
-
-      translated_texts[translated_count] = translated_text
-      translated_count += 1
-
-      if did_touch_end:
-        break
-
-    if spine:
-      spine_path, spine_file = spine
-      if translated_count > 0:
-        spine_file.write_texts(translated_texts, append)
-      context.write_spine_file(spine_path, spine_file)
-
-def _gen_fragments(context: ZipContext):
-  for spine_path in context.search_spine_paths():
-    spine_file = context.read_spine_file(spine_path)
-    for text in spine_file.read_texts():
-      yield Fragment(
-        text=text,
-        start_incision=Incision.IMPOSSIBLE,
-        end_incision=Incision.IMPOSSIBLE,
-      )
-
-def _clean_path(path: Path) -> Path:
-  if path.exists():
-    if path.is_file():
-      path.unlink()
-    elif path.is_dir():
-      rmtree(path, ignore_errors=True)
-  return path
+    llm: LLM,
+    source_path: Path,
+    target_path: Path,
+    target_language: str,
+    user_prompt: str | None = None,
+    max_retries: int = 5,
+    max_group_tokens: int = 1200,
+    on_progress: Callable[[float], None] | None = None,
+) -> None:
+    translator = XMLTranslator(
+        llm=llm,
+        target_language=target_language,
+        user_prompt=user_prompt,
+        ignore_translated_error=False,
+        max_retries=max_retries,
+        max_fill_displaying_errors=10,
+        group_context=XMLGroupContext(
+            encoding=llm.encoding,
+            max_group_tokens=max_group_tokens,
+        ),
+    )
+    with Zip(source_path, target_path) as zip:
+        # Progress distribution: TOC 3%, metadata 2%, chapters 95%
+        TOC_PROGRESS = 0.03
+        METADATA_PROGRESS = 0.02
+        CHAPTERS_PROGRESS = 0.95
+
+        # Count total chapters for progress calculation (lightweight, no content loading)
+        total_chapters = _count_chapters(zip)
+        chapter_progress_step = CHAPTERS_PROGRESS / total_chapters if total_chapters > 0 else 0
+
+        current_progress = 0.0
+
+        # Translate TOC
+        _translate_toc(translator, zip)
+        current_progress += TOC_PROGRESS
+        if on_progress:
+            on_progress(current_progress)
+
+        # Translate metadata
+        _translate_metadata(translator, zip)
+        current_progress += METADATA_PROGRESS
+        if on_progress:
+            on_progress(current_progress)
+
+        # Translate chapters
+        processed_chapters = 0
+        for element, text_segments, (chapter_path, xml, placeholder) in translator.translate_to_text_segments(
+            items=_search_chapter_items(zip),
+        ):
+            submit_text_segments(
+                element=element,
+                text_segments=(
+                    segment
+                    for segment in text_segments
+                    if not any(is_placeholder_tag(e.tag) for e in segment.parent_stack)
+                ),
+            )
+            placeholder.recover()
+            deduplicate_ids_in_element(xml.element)
+            with zip.replace(chapter_path) as target_file:
+                xml.save(target_file)
+
+            # Update progress after each chapter
+            processed_chapters += 1
+            current_progress = TOC_PROGRESS + METADATA_PROGRESS + (processed_chapters * chapter_progress_step)
+            if on_progress:
+                on_progress(current_progress)
+
+
+def _translate_toc(translator: XMLTranslator, zip: Zip):
+    """Translate TOC (Table of Contents) titles."""
+    toc_list = read_toc(zip)
+    if not toc_list:
+        return
+
+    # Collect all titles recursively
+    titles_to_translate: list[str] = []
+
+    def collect_titles(items):
+        for item in items:
+            titles_to_translate.append(item.title)
+            if item.children:
+                collect_titles(item.children)
+
+    collect_titles(toc_list)
+
+    # Create XML elements for translation
+    elements_to_translate = Element("toc")
+    elements_to_translate.extend(_create_text_element(title) for title in titles_to_translate)
+
+    # Translate all titles at once
+    translated_element = translator.translate_to_element(elements_to_translate)
+
+    # Extract translated texts
+    from builtins import zip as builtin_zip
+
+    translated_titles = [
+        plain_text(elem) if elem is not None else original
+        for elem, original in builtin_zip(translated_element, titles_to_translate)
+    ]
+
+    # Fill back translated titles
+    title_index = 0
+
+    def fill_titles(items):
+        nonlocal title_index
+        for item in items:
+            item.title = translated_titles[title_index]
+            title_index += 1
+            if item.children:
+                fill_titles(item.children)
+
+    fill_titles(toc_list)
+
+    # Write back the translated TOC
+    write_toc(zip, toc_list)
+
+
+def _translate_metadata(translator: XMLTranslator, zip: Zip):
+    """Translate metadata fields in OPF file."""
+    opf_path = find_opf_path(zip)
+
+    with zip.read(opf_path) as f:
+        xml = XMLLikeNode(f)
+
+    # Find metadata element
+    metadata_elem = None
+    for child in xml.element:
+        if child.tag.endswith("metadata"):
+            metadata_elem = child
+            break
+
+    if metadata_elem is None:
+        return
+
+    # Collect metadata fields to translate
+    # Skip fields that should not be translated
+    skip_fields = {
+        "language",
+        "identifier",
+        "date",
+        "meta",
+        "contributor",  # Usually technical information
+    }
+
+    fields_to_translate: list[tuple[Element, str]] = []
+
+    for elem in metadata_elem:
+        # Get tag name without namespace
+        tag_name = elem.tag.split("}")[-1] if "}" in elem.tag else elem.tag
+
+        # Check if element has text content and should be translated
+        if elem.text and elem.text.strip() and tag_name not in skip_fields:
+            fields_to_translate.append((elem, elem.text.strip()))
+
+    if not fields_to_translate:
+        return
+
+    # Create XML elements for translation
+    elements_to_translate = Element("metadata")
+    elements_to_translate.extend(_create_text_element(text) for _, text in fields_to_translate)
+
+    # Translate all metadata at once
+    translated_element = translator.translate_to_element(elements_to_translate)
+
+    # Fill back translated texts
+    from builtins import zip as builtin_zip
+
+    for (elem, _), translated_elem in builtin_zip(fields_to_translate, translated_element, strict=True):
+        if translated_elem is not None:
+            translated_text = plain_text(translated_elem)
+            if translated_text:
+                elem.text = translated_text
+
+    # Write back the modified OPF file
+    with zip.replace(opf_path) as f:
+        xml.save(f)
+
+
+def _count_chapters(zip: Zip) -> int:
+    """Count total chapters without loading content (lightweight)."""
+    return sum(1 for _ in search_spine_paths(zip))
+
+
+def _search_chapter_items(zip: Zip):
+    for chapter_path in search_spine_paths(zip):
+        with zip.read(chapter_path) as chapter_file:
+            xml = XMLLikeNode(
+                file=chapter_file,
+                is_html_like=chapter_path.suffix.lower() in (".html", ".htm"),
+            )
+        body_element = find_first(xml.element, "body")
+        if body_element is not None:
+            placeholder = Placeholder(body_element)
+            yield body_element, (chapter_path, xml, placeholder)
+
+
+def _create_text_element(text: str) -> Element:
+    elem = Element("text")
+    elem.text = text
+    return elem

epub_translator/utils.py

@@ -0,0 +1,7 @@
+import re
+
+_WHITESPACE_PATTERN = re.compile(r"\s+")
+
+
+def normalize_whitespace(text: str) -> str:
+    return _WHITESPACE_PATTERN.sub(" ", text)

epub_translator/xml/__init__.py

@@ -1,3 +1,4 @@
-from .encoder import encode, encode_friendly
-from .decoder import decode_friendly
-from .utils import clone
+from .deduplication import *
+from .firendly import *
+from .xml import *
+from .xml_like import *

epub_translator/xml/deduplication.py

@@ -0,0 +1,38 @@
+from xml.etree.ElementTree import Element
+
+from .xml import iter_with_stack
+
+_ID_KEY = "id"
+_SUFFIX = "__translated"
+
+
+def deduplicate_ids_in_element(element: Element) -> Element:
+    seen_ids: set[str] = set()
+    original_id_count: dict[str, int] = {}
+
+    for _, sub_element in iter_with_stack(element):
+        if _ID_KEY not in sub_element.attrib:
+            continue
+        original_id = sub_element.attrib[_ID_KEY]
+
+        if original_id not in seen_ids:
+            seen_ids.add(original_id)
+            original_id_count[original_id] = 1
+        else:
+            original_id_count[original_id] = original_id_count.get(original_id, 1) + 1
+            occurrence = original_id_count[original_id]
+
+            if occurrence == 2:
+                new_id = f"{original_id}{_SUFFIX}"
+            else:
+                new_id = f"{original_id}{_SUFFIX}_{occurrence - 1}"
+
+            counter = occurrence - 1
+            while new_id in seen_ids:
+                counter += 1
+                new_id = f"{original_id}{_SUFFIX}_{counter}"
+
+            sub_element.attrib["id"] = new_id
+            seen_ids.add(new_id)
+
+    return element

epub_translator/xml/firendly/__init__.py

@@ -0,0 +1,2 @@
+from .decoder import decode_friendly
+from .encoder import encode_friendly

epub_translator/xml/firendly/decoder.py

@@ -0,0 +1,75 @@
+from collections.abc import Generator, Iterable
+from xml.etree.ElementTree import Element
+
+from ..xml import clone_element
+from .parser import parse_tags
+from .tag import Tag, TagKind
+from .transform import tag_to_element
+
+
+# why implement XML decoding?
+# https://github.com/oomol-lab/pdf-craft/issues/149
+def decode_friendly(chars: Iterable[str], tags: Iterable[str] | str = ()) -> Generator[Element, None, None]:
+    if isinstance(tags, str):
+        tags = set((tags,))
+    else:
+        tags = set(tags)
+
+    for element in _collect_elements(chars):
+        if element.tag in tags or len(tags) == 0:
+            yield clone_element(element)
+
+
+def _collect_elements(chars: Iterable[str]) -> Generator[Element, None, None]:
+    opening_stack: list[Element] = []
+    last_closed_element: Element | None = None
+
+    for cell in parse_tags(chars):
+        if isinstance(cell, Tag):
+            tag: Tag = cell
+            element = tag_to_element(tag)
+            if tag.kind == TagKind.CLOSING:
+                popped = _pop_element(tag.name, opening_stack)
+                if popped is not None:
+                    yield popped
+                    last_closed_element = popped
+                elif last_closed_element is not None:
+                    _append_to_tail(last_closed_element, tag.proto)
+            else:
+                if opening_stack:
+                    opening_stack[-1].append(element)
+                if tag.kind == TagKind.SELF_CLOSING:
+                    yield element
+                    last_closed_element = element
+                elif tag.kind == TagKind.OPENING:
+                    opening_stack.append(element)
+                    last_closed_element = None
+
+        elif last_closed_element is not None:
+            _append_to_tail(last_closed_element, cell)
+
+        elif opening_stack:
+            opening_stack[-1].text = cell
+
+
+def _append_to_tail(element: Element, text: str) -> None:
+    if element.tail:
+        element.tail += text
+    else:
+        element.tail = text
+
+
+def _pop_element(tag_name: str, opening_stack: list[Element]) -> Element | None:
+    index = -1
+    for i in range(len(opening_stack) - 1, -1, -1):
+        opening_element = opening_stack[i]
+        if tag_name == opening_element.tag:
+            index = i
+            break
+    if index == -1:
+        return None
+
+    popped: Element | None = None
+    for _ in range(len(opening_stack) - index):
+        popped = opening_stack.pop()
+    return popped

epub_translator/xml/firendly/encoder.py

@@ -0,0 +1,84 @@
+from collections.abc import Callable
+from html import escape as escape_html
+from io import StringIO
+from xml.etree.ElementTree import Element
+
+from .parser import parse_tags
+from .tag import Tag, TagKind
+from .transform import element_to_tag
+
+
+# why implement XML encoding?
+# https://github.com/oomol-lab/pdf-craft/issues/149
+def encode_friendly(element: Element, indent: int = 2) -> str:
+    buffer = StringIO()
+    _encode_element(
+        buffer=buffer,
+        element=element,
+        indent=indent,
+        depth=0,
+        escape=_escape_text,
+    )
+    return buffer.getvalue()
+
+
+def _escape_text(text: str) -> str:
+    buffer = StringIO()
+    for cell in parse_tags(text):
+        if isinstance(cell, Tag):
+            cell = escape_html(str(cell))
+        buffer.write(cell)
+    return buffer.getvalue()
+
+
+_TINY_TEXT_LEN = 35
+
+
+def _encode_element(
+    buffer: StringIO,
+    element: Element,
+    indent: int,
+    depth: int,
+    escape: Callable[[str], str],
+) -> None:
+    _write_indent(buffer, indent, depth)
+    if len(element) == 0 and not element.text:
+        tag = element_to_tag(element, TagKind.SELF_CLOSING)
+        buffer.write(str(tag))
+    else:
+        text = (element.text or "").strip()
+        opening_tag = element_to_tag(element, TagKind.OPENING)
+        closing_tag = element_to_tag(element, TagKind.CLOSING)
+        buffer.write(str(opening_tag))
+        is_one_line = len(text) <= _TINY_TEXT_LEN and len(element) == 0 and "\n" not in text
+        if text:
+            if not is_one_line:
+                buffer.write("\n")
+                _write_indent(buffer, indent, depth + 1)
+            buffer.write(escape(text))
+
+        for child in element:
+            buffer.write("\n")
+            _encode_element(
+                buffer=buffer,
+                element=child,
+                indent=indent,
+                depth=depth + 1,
+                escape=escape,
+            )
+            child_tail = (child.tail or "").strip()
+            if child_tail:
+                buffer.write("\n")
+                _write_indent(buffer, indent, depth + 1)
+                buffer.write(escape(child_tail))
+
+        if not is_one_line:
+            buffer.write("\n")
+            _write_indent(buffer, indent, depth)
+
+        buffer.write(str(closing_tag))
+
+
+def _write_indent(buffer: StringIO, indent: int, depth: int) -> None:
+    for _ in range(indent * depth):
+        buffer.write(" ")