Assign formsets to the correct tab

Fixes #255

wagtail_localize/test/migrations/0013_pagewithcustomedithandler_pagewithcustomedithandlerchildobject.py

@@ -0,0 +1,46 @@
+# Generated by Django 3.1.3 on 2020-12-04 16:58
+
+from django.db import migrations, models
+import django.db.models.deletion
+import modelcluster.fields
+import uuid
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('wagtailcore', '0059_apply_collection_ordering'),
+        ('wagtail_localize_test', '0012_auto_20201201_1924'),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name='PageWithCustomEditHandler',
+            fields=[
+                ('page_ptr', models.OneToOneField(auto_created=True, on_delete=django.db.models.deletion.CASCADE, parent_link=True, primary_key=True, serialize=False, to='wagtailcore.page')),
+                ('foo_field', models.TextField()),
+                ('bar_field', models.TextField()),
+                ('baz_field', models.TextField()),
+            ],
+            options={
+                'abstract': False,
+            },
+            bases=('wagtailcore.page',),
+        ),
+        migrations.CreateModel(
+            name='PageWithCustomEditHandlerChildObject',
+            fields=[
+                ('id', models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
+                ('translation_key', models.UUIDField(default=uuid.uuid4, editable=False)),
+                ('sort_order', models.IntegerField(blank=True, editable=False, null=True)),
+                ('field', models.TextField()),
+                ('locale', models.ForeignKey(editable=False, on_delete=django.db.models.deletion.PROTECT, related_name='+', to='wagtailcore.locale')),
+                ('page', modelcluster.fields.ParentalKey(on_delete=django.db.models.deletion.CASCADE, related_name='child_objects', to='wagtail_localize_test.pagewithcustomedithandler')),
+            ],
+            options={
+                'ordering': ['sort_order'],
+                'abstract': False,
+                'unique_together': {('translation_key', 'locale')},
+            },
+        ),
+    ]

wagtail_localize/test/models.py

@@ -2,7 +2,7 @@
 from django.utils.translation import gettext_lazy
 
 from modelcluster.fields import ParentalKey
-from wagtail.admin.edit_handlers import FieldPanel, InlinePanel, PageChooserPanel
+from wagtail.admin.edit_handlers import TabbedInterface, ObjectList, FieldPanel, InlinePanel, PageChooserPanel
 from wagtail.core import blocks
 from wagtail.core.fields import RichTextField, StreamField
 from wagtail.core.models import Page, Orderable, TranslatableMixin
@@ -301,3 +301,35 @@ class TestModelWithInvalidForeignKey(TranslatableMixin, models.Model):
     translatable_fields = [
         TranslatableField('fk'),
     ]
+
+
+class PageWithCustomEditHandler(Page):
+    foo_field = models.TextField()
+    bar_field = models.TextField()
+    baz_field = models.TextField()
+
+    foo_panels = [
+        FieldPanel('foo_field'),
+    ]
+
+    bar_panels = [
+        FieldPanel('bar_field'),
+        FieldPanel('baz_field'),
+    ]
+
+    edit_handler = TabbedInterface([
+        ObjectList(bar_panels, heading="Bar"),
+        ObjectList([InlinePanel('child_objects')], heading="Child objects"),
+        ObjectList(foo_panels, heading="Foo"),
+        ObjectList(Page.content_panels, heading="Content"),
+        ObjectList(Page.promote_panels, heading="Promote"),
+        ObjectList(Page.settings_panels, heading="Settings"),
+    ])
+
+
+class PageWithCustomEditHandlerChildObject(TranslatableMixin, Orderable):
+    page = ParentalKey(PageWithCustomEditHandler, related_name="child_objects")
+    field = models.TextField()
+
+    class Meta(TranslatableMixin.Meta, Orderable.Meta):
+        pass

wagtail_localize/tests/test_edit_translation.py

@@ -24,7 +24,7 @@
 from wagtail.tests.utils import WagtailTestUtils
 
 from wagtail_localize.models import SegmentOverride, String, StringTranslation, Translation, TranslationContext, TranslationLog, TranslationSource, OverridableSegment
-from wagtail_localize.test.models import TestPage, TestSnippet, NonTranslatableSnippet, TestHomePage
+from wagtail_localize.test.models import TestPage, TestSnippet, NonTranslatableSnippet, TestHomePage, PageWithCustomEditHandler, PageWithCustomEditHandlerChildObject
 from wagtail_localize.wagtail_hooks import SNIPPET_RESTART_TRANSLATION_ENABLED
 
 from .utils import assert_permission_denied
@@ -312,6 +312,69 @@ def test_override_types(self):
             ]
         )
 
+    def test_edit_page_translation_with_custom_edit_handler(self):
+        page = self.home_page.add_child(
+            instance=PageWithCustomEditHandler(
+                title="Custom edit handler",
+                foo_field="Foo",
+                bar_field="Bar",
+                baz_field="Baz",
+                child_objects=[
+                    PageWithCustomEditHandlerChildObject(
+                        field="Test 1",
+                    ),
+                    PageWithCustomEditHandlerChildObject(
+                        field="Test 2",
+                    ),
+                ]
+            )
+        )
+
+        source, created = TranslationSource.get_or_create_from_instance(page)
+        translation = Translation.objects.create(
+            source=source,
+            target_locale=self.fr_locale,
+        )
+        translation.save_target()
+        fr_page = page.get_translation(self.fr_locale)
+
+        response = self.client.get(reverse('wagtailadmin_pages:edit', args=[fr_page.id]))
+        self.assertEqual(response.status_code, 200)
+        self.assertTemplateUsed(response, 'wagtail_localize/admin/edit_translation.html')
+
+        # Check props
+        props = json.loads(response.context['props'])
+
+        self.assertEqual(props['tabs'], [
+            {'label': 'Bar', 'slug': 'bar'},
+            {'label': 'Child objects', 'slug': 'child-objects'},
+            {'label': 'Foo', 'slug': 'foo'},
+            {'label': 'Content', 'slug': 'content'},
+            {'label': 'Promote', 'slug': 'promote'},
+            {'label': 'Settings', 'slug': 'settings'},
+        ])
+
+        # Test locations
+        locations = [
+            {
+                'value': segment['source'],
+                'tab': segment['location']['tab'],
+                'field': segment['location']['field'],
+                'model_order': segment['order'],
+                'panel_order': segment['location']['order'],
+            }
+            for segment in props['segments']
+        ]
+        self.assertEqual(locations, [
+            {'value': 'Bar', 'tab': 'bar', 'field': 'Bar field', 'model_order': 4, 'panel_order': 0},
+            {'value': 'Baz', 'tab': 'bar', 'field': 'Baz field', 'model_order': 5, 'panel_order': 1},
+            {'value': 'Test 1', 'tab': 'child-objects', 'field': 'Page with custom edit handler child object', 'model_order': 6, 'panel_order': 2},
+            {'value': 'Test 2', 'tab': 'child-objects', 'field': 'Page with custom edit handler child object', 'model_order': 7, 'panel_order': 2},
+            {'value': 'Foo', 'tab': 'foo', 'field': 'Foo field', 'model_order': 3, 'panel_order': 3},
+            {'value': 'Custom edit handler', 'tab': 'content', 'field': 'Title', 'model_order': 1, 'panel_order': 4},
+            {'value': 'custom-edit-handler', 'tab': 'promote', 'field': 'Slug', 'model_order': 2, 'panel_order': 5}
+        ])
+
     def test_edit_page_translation_with_multi_mode_preview(self):
         # Add some extra preview modes to the page
         previous_preview_modes = TestPage.preview_modes

wagtail_localize/views/edit_translation.py

@@ -142,6 +142,9 @@ def field_tab_mapping(self):
                 for tab_field in tab.required_fields():
                     field_tabs[tab_field] = tab.heading
 
+                for tab_formset in tab.required_formsets().keys():
+                    field_tabs[tab_formset] = tab.heading
+
             return field_tabs
         else:
             return {}
@@ -150,7 +153,7 @@ def get_field_tab(self, field_name):
         if field_name in self.field_tab_mapping:
             return self.field_tab_mapping[field_name]
         else:
-            return self.tabs[0]
+            raise KeyError(f"Cannot find tab for field '{field_name}''")
 
     @cached_property
     def field_ordering_mapping(self):
@@ -167,13 +170,19 @@ def field_ordering_mapping(self):
                     field_orderings[tab_field] = order
                     order += 1
 
+                for tab_formset in tab.required_formsets().keys():
+                    field_orderings[tab_formset] = order
+                    order += 1
+
             return field_orderings
         else:
             return {}
 
     def get_field_order(self, field_name):
         if field_name in self.field_ordering_mapping:
             return self.field_ordering_mapping[field_name]
+        else:
+            raise KeyError(f"Cannot find ordering for field '{field_name}''")
 
     @cached_property
     def field_edit_handler_mapping(self):