makedbmessages.py

from __future__ import annotations

import os
from pathlib import Path

import polib
from django.core.management.base import BaseCommand
from modeltranslation.translator import translator

from instawork import settings


class Command(BaseCommand):
    help = (
        "Dumps the various model fields tagged for translation into models.po files. "
        "This replaces the existing .po file, if any, with a new dump corresponding to "
        "the database contents. Language detection is done via the settings.py language variables."
    )

    def add_arguments(self, parser):
        parser.add_argument(
            "--outdir",
            default=None,
            help="override the output directory. Defaults to the locale directory in settings.py",
        )
        parser.add_argument("--execute", default=False, action="store_true")
        parser.add_argument(
            "--nobackup", default=False, action="store_true", help="Skips backing up of an existing po file"
        )
        parser.add_argument("--verbose", default=False, action="store_true")
        parser.add_argument(
            "--handle-dupes",
            default="raise",
            choices=["raise", "warn"],
            help=(
                "Select how to handle duplicate msgids with different existing translations. "
                "raise: raise when a duplicate is found; warn: Keep original, but warn; "
            ),
        )

    def update_pofile_with_entry(
        self, pofile: polib.POFile, msg_key: tuple[str, str], text_main: str, text_lang: str, handle_dupes: str
    ):
        """Updates a POFile with a single entry, handling the various duplication conditions"""
        try:
            # Add occurrences to an existing entry if one exists
            entry = [e for e in pofile if e.msgid == text_main][0]
            # Override missing translations / Handle conflicts
            if entry.msgstr != text_lang and entry.msgstr and text_lang:
                # A false translation is when the translation is the same as the main language
                # Overwrite existing False translations, if present
                if entry.msgstr == entry.msgid:
                    entry.msgstr = text_lang
                # If this itself is a false translation, skip
                elif text_lang == text_main:
                    pass
                else:
                    msg_conflict = f"Conflict for msgid={entry.msgid}: existing '{entry.msgstr}' != new '{text_lang}'\n"
                    if handle_dupes == "warn":
                        self.stdout.write(self.style.WARNING(msg_conflict))
                    else:
                        raise ValueError(msg_conflict)
            else:
                entry.msgstr = entry.msgstr or text_lang
            entry.occurrences.append(msg_key)
        except IndexError:
            entry = polib.POEntry(
                occurrences=[msg_key],
                msgid=text_main,
                msgstr=text_lang,
            )
            pofile.append(entry)

    def generate_language_pofile(
        self, lang_main: str, lang_target: str, handle_dupes="raise", **options
    ) -> polib.POFile:
        """
        Generate the POFile object for a given language's dump. Only creates an entry
        ### Args
        - `lang_main`: 2 Character Language code for the reference language
        - `lang_target`: 2 Character Language code for the translation target
        ### Returns
        - `POFile`: In-memory `POFile` For the specified language pair
        """
        pofile = polib.POFile()
        for model in translator.get_registered_models():
            fields = translator.get_options_for_model(model).get_field_names()
            self.stdout.write(f"{model}: {fields}")
            model_objs = model.objects.all()
            for field in fields:
                # Get data for all fields of interest, if present
                # Save returns an error if nonetype passed
                field_data = [
                    (
                        model_obj.pk,
                        getattr(model_obj, f"{field}_{lang_target}") or "",
                        getattr(model_obj, f"{field}_{lang_main}"),
                    )
                    for model_obj in model_objs
                    # We do not populate the .PO file when no primary data exists.
                    if getattr(model_obj, f"{field}_{lang_main}")
                ]
                for pk, text_lang, text_main in field_data:
                    # Cast pk to string here, or it fails silently
                    msg_key = (f"{model._meta.app_label}.{model.__name__}.{field}", str(pk))

                    if options["verbose"]:
                        self.stdout.write(f"{msg_key}: {text_main} -> {text_lang}\n")
                    self.update_pofile_with_entry(pofile, msg_key, text_main, text_lang, handle_dupes=handle_dupes)

        return pofile

    def export_model_data_po(
        self,
        lang_targets: None | list = None,
        execute: bool = True,
        verbose: bool = False,
        outdir: None | str | Path = None,
        handle_dupes: str = "raise",
        **options,
    ) -> None:
        """
        For each language creates a models.po file in which the
        msgid is the key corresponding to the item, and the msgstr
        the translation, if any
        ### Args
        - `execute`: Apply changes, writing to disk
        - `verbose`: Display all exported messages to stdout
        - `outdir`: Override the output directory. defaults to the locale dir otherwise
        - `lang_targets`: Override target languages. Defaults to all languages in settings.py, main or otherwise.
        - `handle_dupes`: How to handle duplicates. "raise" raises if a dupe is found. "warn" prints a warning.
        """
        if not outdir:
            outdir = settings.LOCALE_PATHS[0]

        # Base vs translation language configuration
        lang_main = settings.LANGUAGE_CODE.split("-")[0]
        # This is duplication safe, and lang_main missing from languages safe
        if not lang_targets:
            lang_targets = [lang_code for lang_code, lang_name in settings.LANGUAGES if lang_code != lang_main] + [
                lang_main
            ]

        self.stdout.write("Dumping model translation fields\n")
        self.stdout.write(f"\tbase lang: {lang_main}, target langs: {lang_targets}\n")
        for lang_target in lang_targets:
            os.makedirs(os.path.join(outdir, lang_target, "LC_MESSAGES"), exist_ok=True)
            pofile = self.generate_language_pofile(lang_main, lang_target, verbose=verbose, handle_dupes=handle_dupes)
            if not execute:
                continue
            path = os.path.join(outdir, lang_target, "LC_MESSAGES", "models.po")
            pofile.save(path)

    def handle(self, *args, **options):
        self.export_model_data_po(**options)