Consistently and safely set locale params in generic views

wagtail/admin/views/generic/mixins.py

@@ -18,7 +18,7 @@ from wagtail import hooks
 from wagtail.admin import messages
 from wagtail.admin.templatetags.wagtailadmin_tags import user_display_name
 from wagtail.admin.ui.tables import TitleColumn
-from wagtail.admin.utils import get_latest_str
+from wagtail.admin.utils import get_latest_str, set_query_params
 from wagtail.locks import BasicLock, ScheduledForPublishLock, WorkflowLock
 from wagtail.log_actions import log
 from wagtail.log_actions import registry as log_registry
@@ -135,6 +135,11 @@ class LocaleMixin:
         context["translations"] = self.translations
         return context
 
+    def _set_locale_query_param(self, url, locale=None):
+        if not (locale := locale or self.locale):
+            return url
+        return set_query_params(url, {"locale": locale.language_code})
+
 
 class PanelMixin:
     panel = None

wagtail/admin/views/generic/models.py

@@ -398,7 +398,7 @@ class IndexView(
         return [
             {
                 "locale": locale,
-                "url": index_url + "?locale=" + locale.language_code,
+                "url": self._set_locale_query_param(index_url, locale),
             }
             for locale in Locale.objects.all().exclude(id=self.locale.id)
         ]
@@ -571,7 +571,7 @@ class CreateView(
                 "Subclasses of wagtail.admin.views.generic.models.CreateView must provide an "
                 "add_url_name attribute or a get_add_url method"
             )
-        return reverse(self.add_url_name)
+        return self._set_locale_query_param(reverse(self.add_url_name))
 
     def get_edit_url(self):
         if not self.edit_url_name:
@@ -587,7 +587,7 @@ class CreateView(
                 "Subclasses of wagtail.admin.views.generic.models.CreateView must provide an "
                 "index_url_name attribute or a get_success_url method"
             )
-        return reverse(self.index_url_name)
+        return self._set_locale_query_param(reverse(self.index_url_name))
 
     def get_success_message(self, instance):
         if self.success_message is None:
@@ -609,10 +609,11 @@ class CreateView(
         return context
 
     def get_translations(self):
+        add_url = self.get_add_url()
         return [
             {
                 "locale": locale,
-                "url": self.get_add_url() + "?locale=" + locale.language_code,
+                "url": self._set_locale_query_param(add_url, locale),
             }
             for locale in Locale.objects.all().exclude(id=self.locale.id)
         ]

wagtail/snippets/tests/test_snippets.py

@@ -920,6 +920,32 @@ class TestLocaleSelectorOnCreate(WagtailTestUtils, TestCase):
 
         self.assertContains(response, "Switch locales")
 
+        switch_to_french_url = (
+            reverse("wagtailsnippets_snippetstests_translatablesnippet:add")
+            + "?locale=fr"
+        )
+        self.assertContains(
+            response,
+            f'<a href="{switch_to_french_url}" lang="fr">',
+        )
+
+    def test_locale_selector_with_existing_locale(self):
+        response = self.client.get(
+            reverse("wagtailsnippets_snippetstests_translatablesnippet:add")
+            + "?locale=fr"
+        )
+
+        self.assertContains(response, "Switch locales")
+
+        switch_to_english_url = (
+            reverse("wagtailsnippets_snippetstests_translatablesnippet:add")
+            + "?locale=en"
+        )
+        self.assertContains(
+            response,
+            f'<a href="{switch_to_english_url}" lang="en">',
+        )
+
     @override_settings(WAGTAIL_I18N_ENABLED=False)
     def test_locale_selector_not_present_when_i18n_disabled(self):
         response = self.client.get(
@@ -928,11 +954,29 @@ class TestLocaleSelectorOnCreate(WagtailTestUtils, TestCase):
 
         self.assertNotContains(response, "Switch locales")
 
+        switch_to_french_url = (
+            reverse("wagtailsnippets_snippetstests_translatablesnippet:add")
+            + "?locale=fr"
+        )
+        self.assertNotContains(
+            response,
+            f'<a href="{switch_to_french_url}" lang="fr">',
+        )
+
     def test_locale_selector_not_present_on_non_translatable_snippet(self):
         response = self.client.get(reverse("wagtailsnippets_tests_advert:add"))
 
         self.assertNotContains(response, "Switch locales")
 
+        switch_to_french_url = (
+            reverse("wagtailsnippets_snippetstests_translatablesnippet:add")
+            + "?locale=fr"
+        )
+        self.assertNotContains(
+            response,
+            f'<a href="{switch_to_french_url}" lang="fr">',
+        )
+
 
 class TestCreateDraftStateSnippet(WagtailTestUtils, TestCase):
     STATUS_TOGGLE_BADGE_REGEX = (

wagtail/snippets/views/snippets.py

@@ -39,7 +39,6 @@ from wagtail.admin.viewsets.model import ModelViewSet, ModelViewSetGroup
 from wagtail.admin.widgets.button import BaseDropdownMenuButton, ButtonWithDropdown
 from wagtail.models import (
     DraftStateMixin,
-    Locale,
     LockableMixin,
     PreviewableMixin,
     RevisionMixin,
@@ -245,23 +244,6 @@ class CreateView(generic.CreateEditViewOptionalFeaturesMixin, generic.CreateView
     def run_after_hook(self):
         return self.run_hook("after_create_snippet", self.request, self.object)
 
-    def get_add_url(self):
-        url = reverse(self.add_url_name)
-        if self.locale:
-            url += "?locale=" + self.locale.language_code
-        return url
-
-    def get_success_url(self):
-        if self.draftstate_enabled and self.action != "publish":
-            return super().get_success_url()
-
-        # Make sure the redirect to the listing view uses the correct locale
-        urlquery = ""
-        if self.locale and self.object.locale is not Locale.get_default():
-            urlquery = "?locale=" + self.object.locale.language_code
-
-        return reverse(self.index_url_name) + urlquery
-
     def _get_action_menu(self):
         return SnippetActionMenu(self.request, view=self.view_name, model=self.model)