epub-translator 0.0.6__py3-none-any.whl → 0.1.0__py3-none-any.whl

epub_translator/serial/chunk.py

@@ -0,0 +1,52 @@
+from collections.abc import Generator, Iterable
+from dataclasses import dataclass
+from typing import Generic
+
+from resource_segmentation import Resource, Segment, split
+
+from .segment import ST
+
+_INCISION = 0
+
+
+@dataclass
+class Chunk(Generic[ST]):
+    head_remain_tokens: int
+    tail_remain_tokens: int
+    head: list[ST]
+    body: list[ST]
+    tail: list[ST]
+
+
+def split_into_chunks(segments: Iterable[ST], max_group_tokens: int) -> Generator[Chunk[ST], None, None]:
+    for group in split(
+        max_segment_count=max_group_tokens,
+        gap_rate=0.07,
+        tail_rate=0.5,
+        border_incision=_INCISION,
+        resources=(
+            Resource(
+                count=segment.tokens,
+                start_incision=_INCISION,
+                end_incision=_INCISION,
+                payload=segment,
+            )
+            for segment in segments
+        ),
+    ):
+        yield Chunk(
+            head_remain_tokens=group.head_remain_count,
+            tail_remain_tokens=group.tail_remain_count,
+            head=list(_expand_payloads(group.head)),
+            body=list(_expand_payloads(group.body)),
+            tail=list(_expand_payloads(group.tail)),
+        )
+
+
+def _expand_payloads(target: list[Resource[ST] | Segment[ST]]) -> Generator[ST, None, None]:
+    for item in target:
+        if isinstance(item, Resource):
+            yield item.payload
+        elif isinstance(item, Segment):
+            for resource in item.resources:
+                yield resource.payload

epub_translator/serial/segment.py

@@ -0,0 +1,17 @@
+from typing import Generic, Protocol, Self, TypeVar, runtime_checkable
+
+S = TypeVar("S", covariant=True)
+T = TypeVar("T")
+ST = TypeVar("ST", bound="Segment")
+
+
+@runtime_checkable
+class Segment(Protocol, Generic[S]):
+    @property
+    def tokens(self) -> int: ...
+
+    @property
+    def payload(self) -> S: ...
+
+    def truncate_after_head(self, remain_tokens: int) -> Self: ...
+    def truncate_before_tail(self, remain_tokens: int) -> Self: ...

epub_translator/serial/splitter.py

@@ -0,0 +1,50 @@
+from collections.abc import Callable, Generator, Iterable
+
+from .chunk import split_into_chunks
+from .segment import ST, T
+
+
+def split(
+    segments: Iterable[ST],
+    transform: Callable[[list[ST]], list[T]],
+    max_group_tokens: int,
+) -> Generator[T, None, None]:
+    for group in split_into_chunks(segments, max_group_tokens):
+        head = list(
+            _truncate_extra_content(
+                segments=group.head,
+                remain_left=False,
+                remain_tokens=group.head_remain_tokens,
+            )
+        )
+        tail = list(
+            _truncate_extra_content(
+                segments=group.tail,
+                remain_left=True,
+                remain_tokens=group.tail_remain_tokens,
+            )
+        )
+        transformed = transform(head + group.body + tail)
+
+        if len(tail) > 0:  # 避免 target[N:-0] 切片错误
+            yield from transformed[len(head) : -len(tail)]
+        else:
+            yield from transformed[len(head) :]
+
+
+def _truncate_extra_content(segments: list[ST], remain_left: bool, remain_tokens: int):
+    tokens_list: list[int] = [segment.tokens for segment in segments]
+    segments = list(segments)
+    for tokens in tokens_list if remain_left else reversed(tokens_list):
+        if remain_tokens <= 0:
+            break
+        next_segment = segments.pop(0) if remain_left else segments.pop()
+        if remain_tokens < tokens:
+            if remain_left:
+                next_segment = next_segment.truncate_after_head(remain_tokens)
+            else:
+                next_segment = next_segment.truncate_before_tail(remain_tokens)
+            remain_tokens = 0
+        else:
+            remain_tokens -= tokens
+        yield next_segment

epub_translator/template.py

@@ -1,50 +1,52 @@
 import re
-
-from typing import Tuple, Callable
+from collections.abc import Callable
 from pathlib import Path
-from jinja2 import select_autoescape, Environment, BaseLoader, TemplateNotFound
+
+from jinja2 import BaseLoader, Environment, TemplateNotFound, select_autoescape
 
 
 def create_env(dir_path: Path) -> Environment:
-  return Environment(
-    loader=_DSLoader(dir_path),
-    autoescape=select_autoescape(),
-    trim_blocks=True,
-    keep_trailing_newline=True,
-  )
+    return Environment(
+        loader=_DSLoader(dir_path),
+        autoescape=select_autoescape(),
+        trim_blocks=True,
+        keep_trailing_newline=True,
+    )
+
+
+_LoaderResult = tuple[str, str | None, Callable[[], bool] | None]
 
-_LoaderResult = Tuple[str, str | None, Callable[[], bool] | None]
 
 class _DSLoader(BaseLoader):
-  def __init__(self, dir_path: Path):
-    super().__init__()
-    self._dir_path: Path = dir_path
+    def __init__(self, dir_path: Path):
+        super().__init__()
+        self._dir_path: Path = dir_path
 
-  def get_source(self, _: Environment, template: str) -> _LoaderResult:
-    template = self._norm_template(template)
-    target_path = (self._dir_path / template).resolve()
+    def get_source(self, environment: Environment, template: str) -> _LoaderResult:
+        template = self._norm_template(template)
+        target_path = (self._dir_path / template).resolve()
 
-    if not target_path.exists():
-      raise TemplateNotFound(f"cannot find {template}")
+        if not target_path.exists():
+            raise TemplateNotFound(f"cannot find {template}")
 
-    return self._get_source_with_path(target_path)
+        return self._get_source_with_path(target_path)
 
-  def _norm_template(self, template: str) -> str:
-    if bool(re.match(r"^\.+/", template)):
-      raise TemplateNotFound(f"invalid path {template}")
+    def _norm_template(self, template: str) -> str:
+        if bool(re.match(r"^\.+/", template)):
+            raise TemplateNotFound(f"invalid path {template}")
 
-    template = re.sub(r"^/", "", template)
-    template = re.sub(r"\.jinja$", "", template, flags=re.IGNORECASE)
-    template = f"{template}.jinja"
+        template = re.sub(r"^/", "", template)
+        template = re.sub(r"\.jinja$", "", template, flags=re.IGNORECASE)
+        template = f"{template}.jinja"
 
-    return template
+        return template
 
-  def _get_source_with_path(self, path: Path) -> _LoaderResult:
-    mtime = path.stat().st_mtime
-    with open(path, "r", encoding="utf-8") as f:
-      source = f.read()
+    def _get_source_with_path(self, path: Path) -> _LoaderResult:
+        mtime = path.stat().st_mtime
+        with open(path, encoding="utf-8") as f:
+            source = f.read()
 
-    def is_updated() -> bool:
-      return mtime == path.stat().st_mtime
+        def is_updated() -> bool:
+            return mtime == path.stat().st_mtime
 
-    return source, path, is_updated
+        return source, str(path), is_updated

epub_translator/translator.py

@@ -1,174 +1,211 @@
-from os import PathLike
+from collections.abc import Callable
 from pathlib import Path
-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
 
 
 def translate(
-      llm: LLM,
-      source_path: PathLike,
-      translated_path: PathLike,
-      target_language: Language,
-      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,
-    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,
-        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._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
-
-    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)
-        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)
-      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, is_html_like=True)
+
+            # 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(chapter_file)
+        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