translator.py

import asyncio
import json
import os
import random
import time
from typing import Generator, Tuple

from docx import Document
from googletrans import Translator

# Language codes used across the app
LANG_SOURCE_CODES = ["auto", "zh-cn", "en", "fr", "es"]
LANG_TARGET_CODES = ["zh-cn", "en", "fr", "es"]


def load_or_initialize_state(state_file: str, num_paragraphs: int) -> dict:
    if os.path.exists(state_file):
        with open(state_file, "r", encoding="utf-8") as f:
            return json.load(f)
    return {
        "translated_count": 0,
        "paragraphs": [],
        "order": [],
        "translated_paragraphs": [None] * num_paragraphs,
    }


def save_state(state_file: str, state: dict) -> None:
    with open(state_file, "w", encoding="utf-8") as f:
        json.dump(state, f, ensure_ascii=False, indent=4)


class DocumentTranslator:
    """Translate a Word document paragraph by paragraph with resumable state."""

    def __init__(
        self,
        input_file: str,
        output_file: str,
        state_file: str,
        src_lang: str = "auto",
        dest_lang: str = "zh-cn",
    ) -> None:
        self.input_file = input_file
        self.output_file = output_file
        self.state_file = state_file
        self.src_lang = src_lang
        self.dest_lang = dest_lang

    def translate(self) -> Generator[Tuple[int, dict], None, None]:
        """
        Run the translation and yield (progress_percent, status_payload).

        Progress updates can be consumed by a GUI thread without blocking the UI.
        """
        doc = Document(self.input_file)
        paragraphs = [(i, para.text) for i, para in enumerate(doc.paragraphs)]

        state = load_or_initialize_state(self.state_file, len(paragraphs))

        if not state["paragraphs"]:
            state["paragraphs"] = paragraphs
            state["order"] = list(range(len(paragraphs)))
            random.shuffle(state["order"])
            save_state(self.state_file, state)

        translator = Translator()
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)

        try:
            total = len(state["order"])
            for i in range(state["translated_count"], total):
                original_index = state["order"][i]
                index, text = state["paragraphs"][original_index]

                if not text.strip():
                    state["translated_paragraphs"][index] = text
                    status = {
                        "event": "skip_empty",
                        "index": i + 1,
                        "total": total,
                    }
                else:
                    result = translator.translate(
                        text, src=self.src_lang, dest=self.dest_lang
                    )
                    # googletrans 3.4.0+ returns a coroutine; handle both sync/async
                    if asyncio.iscoroutine(result):
                        result = loop.run_until_complete(result)
                    translated_text = result.text
                    state["translated_paragraphs"][index] = translated_text
                    status = {
                        "event": "translated",
                        "index": i + 1,
                        "total": total,
                        "src": self.src_lang,
                        "dest": self.dest_lang,
                    }
                    time.sleep(random.uniform(3, 5))

                state["translated_count"] = i + 1
                save_state(self.state_file, state)
                progress = int(((i + 1) / total) * 100)
                yield progress, status

            translated_doc = Document()
            for para in state["translated_paragraphs"]:
                if para:
                    translated_doc.add_paragraph(para)

            translated_doc.save(self.output_file)
            yield 100, {
                "event": "completed",
                "output": self.output_file,
                "src": self.src_lang,
                "dest": self.dest_lang,
            }
        finally:
            # clean up async client
            try:
                loop.run_until_complete(translator.client.aclose())
            except Exception:
                pass
            loop.close()

            # clean state only when translation completes
            if (
                os.path.exists(self.state_file)
                and state.get("translated_count") == len(state.get("order", []))
            ):
                try:
                    os.remove(self.state_file)
                except OSError:
                    pass