Add Management Command for Copying and Translating Pages Across Locales

wagtail_localize/management/commands/translate_all_pages.py

@@ -0,0 +1,191 @@
+import logging
+
+from typing import Callable, Optional
+
+from django.apps import apps
+from django.contrib.auth import get_user_model
+from django.core.management.base import BaseCommand
+from django.db import models
+from wagtail.models import Locale, Page
+
+from wagtail_localize.machine_translators import get_machine_translator
+from wagtail_localize.models import Translation, TranslationSource
+from wagtail_localize.operations import translate_object
+from wagtail_localize.views import edit_translation
+
+
+logger = logging.getLogger(__name__)
+User = get_user_model()
+
+
+class Command(BaseCommand):
+    def add_arguments(self, parser) -> None:
+        parser.add_argument(
+            "language_code", type=str, help="Source language code to copy pages from"
+        )
+        parser.add_argument(
+            "--exclude",
+            nargs="*",
+            type=str,
+            help="List of model names to exclude from translation",
+        )
+        parser.add_argument(
+            "--dry-run",
+            action="store_true",
+            help="Run without making any actual changes",
+        )
+
+    def _get_exclude_models(self, model_names: list[str]) -> list[type[models.Model]]:
+        exclude_models = []
+        for model_name in model_names:
+            found_model = None
+            for model in apps.get_models():
+                if model.__name__ == model_name:
+                    found_model = model
+                    break
+
+            if found_model:
+                exclude_models.append(found_model)
+                self.stdout.write(f"Excluded model: {model_name}")
+            else:
+                self.stderr.write(
+                    f"Model '{model_name}' not found in any installed app."
+                )
+        return exclude_models
+
+    def _get_admin_user(self) -> Optional[User]:
+        admin_user = User.objects.filter(is_superuser=True).first()
+        if not admin_user:
+            self.stderr.write("Error: No superuser found in the system.")
+        return admin_user
+
+    def _handle_page_operation(
+        self, page: Page, operation_func: Callable, error_msg: str, dry_run: bool
+    ) -> bool:
+        if dry_run:
+            self.stdout.write(f"Would process page: {page.title} (ID: {page.id})")
+            return True
+
+        try:
+            operation_func()
+            self.stdout.write(f"Processed page: {page.title} (ID: {page.id})")
+            return True
+        except Exception as e:
+            logger.error(f"{error_msg}: {page.title} (ID: {page.id}): {e}")
+            return False
+
+    def _should_process_page(
+        self, page: Page, exclude_models: list[type[models.Model]]
+    ) -> bool:
+        return not (page.is_root() or isinstance(page.specific, tuple(exclude_models)))
+
+    def _copy_page(
+        self, page: Page, admin_user: User, locales: models.QuerySet
+    ) -> bool:
+        translate_object(page, locales)
+        translation_source, created = TranslationSource.get_or_create_from_instance(
+            page.specific
+        )
+        translation_source.create_or_update_translation(
+            locale=page.locale, user=admin_user, publish=True, fallback=True
+        )
+        translation_source.update_from_db()
+        return True
+
+    def copy_pages(
+        self,
+        admin_user: User,
+        pages_locale_language: models.QuerySet,
+        exclude_models: list[type[models.Model]],
+        dry_run: bool,
+    ) -> None:
+        locales = Locale.objects.all()
+
+        for page in pages_locale_language:
+            if not self._should_process_page(page, exclude_models):
+                continue
+
+            self._handle_page_operation(
+                page,
+                lambda page: self._copy_page(page, admin_user, locales),
+                "Error processing page",
+                dry_run,
+            )
+
+    def _translate_page(self, page: Page, admin_user: User) -> bool:
+        machine_translator = get_machine_translator()
+        translation_source, _ = TranslationSource.update_or_create_from_instance(
+            page.specific
+        )
+
+        translation = Translation.objects.filter(
+            source__object_id=page.translation_key,
+            target_locale_id=page.locale_id,
+            enabled=True,
+        ).first()
+
+        if not translation:
+            return False
+
+        if edit_translation.apply_machine_translation(
+            translation.id, admin_user, machine_translator
+        ):
+            translation_source.create_or_update_translation(
+                locale=page.locale, user=admin_user, publish=True, fallback=True
+            )
+            return True
+        return False
+
+    def translate_pages(
+        self,
+        pages_to_translate: models.QuerySet,
+        excluded_models: list[type[models.Model]],
+        admin_user: User,
+        dry_run: bool,
+    ) -> None:
+        for page in pages_to_translate:
+            if not self._should_process_page(page, excluded_models):
+                continue
+
+            self._handle_page_operation(
+                page,
+                lambda page: self._translate_page(page, admin_user),
+                "Error processing translation",
+                dry_run,
+            )
+
+    def handle(self, *args, **options) -> None:
+        try:
+            Locale.objects.get(language_code=options["language_code"])
+        except Locale.DoesNotExist:
+            self.stderr.write(
+                f"Error: Locale with language code {options['language_code']} not found"
+            )
+            return
+
+        admin_user = self._get_admin_user()
+        if not admin_user:
+            return
+
+        exclude_models = self._get_exclude_models(options["exclude"] or [])
+        pages_locale_language = Page.objects.filter(
+            locale__language_code=options["language_code"]
+        )
+
+        if options["dry_run"]:
+            self.stdout.write("Running in dry-run mode - no changes will be made")
+
+        self.stdout.write("Starting page copy process...")
+        self.copy_pages(
+            admin_user, pages_locale_language, exclude_models, options["dry_run"]
+        )
+
+        pages_to_translate = Page.objects.exclude(
+            locale__language_code=options["language_code"]
+        )
+        self.stdout.write("Starting translation process...")
+        self.translate_pages(
+            pages_to_translate, exclude_models, admin_user, options["dry_run"]
+        )
+
+        self.stdout.write(self.style.SUCCESS("Translation process completed"))

wagtail_localize/tests/test_translate_all_pages.py

@@ -0,0 +1,76 @@
+from io import StringIO
+
+from django.contrib.auth import get_user_model
+from django.core.management import call_command
+from django.test import TestCase
+from wagtail.models import Locale, Page
+
+from wagtail_localize.test.models import TestPage
+
+
+User = get_user_model()
+
+
+class TranslateAllPagesCommandTest(TestCase):
+    def setUp(self):
+        # Set up locales
+        self.default_locale = Locale.objects.get(language_code="en")
+        Locale.objects.create(language_code="fr")
+        Locale.objects.create(language_code="es")
+        User.objects.create(is_superuser=True, is_staff=True, username="admin")
+
+        self.page = Page.objects.get(id=2)
+
+    def test_translate_all_pages_dry_run(self):
+        # Test dry-run mode to ensure no changes are made
+        call_command("translate_all_pages", "en", "--dry-run")
+
+        # Assert no new pages are created
+        self.assertEqual(Page.objects.exclude(locale=self.default_locale).count(), 0)
+
+    def test_translate_all_pages_excludes_models(self):
+        # Test exclude_models argument
+        call_command("translate_all_pages", "en", "--exclude", "Page")
+
+        # Assert the model is excluded from translation
+        self.assertEqual(Page.objects.exclude(locale=self.default_locale).count(), 0)
+
+    def test_translate_all_pages_translation_process(self):
+        # Test successful translation
+        call_command("translate_all_pages", "en")
+
+        # Assert translations are created for target locale
+        translated_pages = Page.objects.exclude(locale=self.default_locale)
+        self.assertEqual(translated_pages.count(), 2)
+
+    def test_translate_all_pages_dry_run_with_exclude_models(self):
+        call_command("translate_all_pages", "en", "--dry-run", "--exclude", "Page")
+        self.assertEqual(Page.objects.exclude(locale=self.default_locale).count(), 0)
+
+    def test_translate_all_pages_excludes_multiple_models(self):
+        call_command("translate_all_pages", "en", "--exclude", "Page", "TestPage")
+
+        # Assert no translated pages
+        self.assertEqual(Page.objects.exclude(locale=self.default_locale).count(), 0)
+        self.assertEqual(
+            TestPage.objects.exclude(locale=self.default_locale).count(), 0
+        )
+
+    def test_translate_all_pages_invalid_exclude_model(self):
+        out = StringIO()
+        err = StringIO()
+
+        call_command(
+            "translate_all_pages",
+            "en",
+            "--exclude",
+            "InvalidModel",
+            verbosity=3,
+            stdout=out,
+            stderr=err,
+        )
+
+        self.assertIn(
+            "Model 'InvalidModel' not found in any installed app", err.getvalue()
+        )
+        self.assertEqual(Page.objects.exclude(locale=self.default_locale).count(), 2)