translate_json.py

# Copyright 2023 DeepL SE (https://www.deepl.com)
# Use of this source code is governed by an MIT
# license that can be found in the LICENSE file.

import deepl
import json
import logging
from typing import Any, Callable, List, Tuple


def batch_translate(
    input_handler_pairs: List[Tuple[str, Callable[[deepl.TextResult], None]]],
    target_lang: str,
    translator: deepl.Translator,
    **kwargs,
) -> None:
    """
    Takes a list of translation inputs and result handlers, translates all
    inputs and calls each handler with the translation results.
    """
    inputs = [input_str for input_str, _ in input_handler_pairs]
    handlers = [handler for _, handler in input_handler_pairs]
    results = translator.translate_text(
        inputs, target_lang=target_lang, **kwargs
    )
    for result, handler in zip(results, handlers):
        handler(result)


def parse_json_for_translation(
    obj: Any,
    translation_candidates: List[
        Tuple[str, Callable[[deepl.TextResult], None]]
    ],
):
    """
    Steps into the given JSON object and adds all strings to
    translation candidates list
    """

    if type(obj) is dict:
        keys = obj.keys()
    elif type(obj) is list:
        keys = range(len(obj))
    else:
        return

    for key in keys:
        if type(obj[key]) is str:
            current_obj = obj
            current_key = key
            assign = lambda val: current_obj.__setitem__(current_key, val.text)
            translation_candidates.append((obj[key], assign))
        else:
            parse_json_for_translation(obj[key], translation_candidates)


def translate_json(
    json_input: str,
    target_lang: str,
    translator: deepl.Translator,
    **kwargs,
) -> str:
    """
    Translates given JSON input using DeepL Translator.

    Most of the arguments of the translate_text function are supported,
    source_lang, target_lang, glossary_id, formality, etc.

    :param json_input: JSON input to be translated.
    :param target_lang: language code to translate template into, for example
        "DE", "EN-US", "FR".
    :param translator: deepl.Translator to use for translation.
    :return: Translated JSON.
    """

    logger = logging.getLogger("deepl")
    obj = json.loads(json_input)
    # Wrap the JSON object in an array, in case the input is a string
    obj = [obj]

    # Find all text in the JSON that is to be translated
    translation_candidates = []
    parse_json_for_translation(obj, translation_candidates)
    logger.info(
        f"Found {len(translation_candidates)} strings to be translated"
    )

    # Translate all texts
    batch_translate(translation_candidates, target_lang, translator, **kwargs)
    logger.info("Translation complete")

    # Unwrap the dummy array and convert to JSON
    return json.dumps(obj[0])