✨ Context records relations: remove form from the wizard to put in a specific form

1 files changed, 49 insertions, 31 deletions

diff --git a/archaeological_context_records/forms.py b/archaeological_context_records/forms.py
index 55f27ddcd..8d41a7888 100644
--- a/archaeological_context_records/forms.py
+++ b/archaeological_context_records/forms.py
@@ -21,6 +21,7 @@
 Context records forms definitions
 """
 from collections import OrderedDict
+from copy import copy
 from itertools import groupby
 
 from bootstrap_datepicker.widgets import DateField
@@ -464,37 +465,6 @@ DatingFormSet.form_admin_name = _("Context record - 030 - Dating")
 DatingFormSet.form_slug = "contextrecord-030-datings"
 
 
-class RecordRelationsForm(OpeRecordRelationsForm):
-    current_model = models.RelationType
-    current_related_model = models.ContextRecord
-    associated_models = {
-        "right_record": models.ContextRecord,
-        "relation_type": models.RelationType,
-    }
-    right_record = forms.ChoiceField(
-        label=_("Context record"), choices=[], required=False
-    )
-
-    def __init__(self, *args, **kwargs):
-        crs = None
-        if "data" in kwargs and "CONTEXT_RECORDS" in kwargs["data"]:
-            crs = kwargs["data"]["CONTEXT_RECORDS"]
-        super(RecordRelationsForm, self).__init__(*args, **kwargs)
-        self.fields["relation_type"].choices = models.RelationType.get_types(
-            initial=self.init_data.get("relation_type")
-        )
-        if crs:
-            self.fields["right_record"].choices = [("", "-" * 2)] + crs
-
-
-RecordRelationsFormSet = formset_factory(
-    RecordRelationsForm, can_delete=True, formset=RecordRelationsFormSetBase
-)
-RecordRelationsFormSet.form_label = _("Relations")
-RecordRelationsFormSet.form_admin_name = _("Context record - 050 - Relations")
-RecordRelationsFormSet.form_slug = "contextrecord-050-recordrelations"
-
-
 class RecordFormInterpretation(CustomForm, ManageOldType):
     HEADERS = {}
     form_label = _("Interpretation")
@@ -554,6 +524,54 @@ class RecordDeletionForm(FinalForm):
     confirm_end_msg = _("Would you like to delete this context record?")
 
 
+class RecordRelationsForm(OpeRecordRelationsForm):
+    current_model = models.RelationType
+    current_related_model = models.ContextRecord
+    associated_models = {
+        "right_record": models.ContextRecord,
+        "relation_type": models.RelationType,
+    }
+    ERROR_MISSING = _("You should select a context record and a relation type.")
+
+    right_record = forms.ChoiceField(
+        label=_("Context record"), choices=[], required=False
+    )
+
+    def __init__(self, *args, **kwargs):
+        crs = None
+        if "data" in kwargs and "CONTEXT_RECORDS" in kwargs["data"]:
+            kwargs["data"] = copy(kwargs["data"])
+            crs = kwargs["data"].pop("CONTEXT_RECORDS")
+            # clean data if not "real" data
+            prefix_value = kwargs['prefix'] + '-relation_type'
+            if not [k for k in kwargs['data'].keys()
+                    if k.startswith(prefix_value) and kwargs['data'][k]]:
+                kwargs.pop('data')
+                if 'files' in kwargs:
+                    kwargs.pop('files')
+        initial = kwargs.get("initial", {})
+        if initial and initial.get("right_record", None):
+            if initial["right_record"] not in [cr_id for cr_id, cr_lbl in crs]:
+                try:
+                    crs.append(
+                        (initial["right_record"],
+                         str(models.ContextRecord.objects.get(pk=initial["right_record"])))
+                    )
+                except models.ContextRecord.DoesNotExist:
+                    pass
+        super().__init__(*args, **kwargs)
+        if crs:
+            self.fields["right_record"].choices = [("", "-" * 2)] + crs
+
+
+RecordRelationsFormSet = formset_factory(
+    RecordRelationsForm, can_delete=True, formset=RecordRelationsFormSetBase
+)
+RecordRelationsFormSet.form_label = _("Relations")
+RecordRelationsFormSet.form_admin_name = _("Context record - 050 - Relations")
+RecordRelationsFormSet.form_slug = "contextrecord-050-recordrelations"
+
+
 class QAOperationCR(IshtarForm):
     town = forms.ChoiceField(label=_("Town"), choices=[])
     archaeological_site = forms.ChoiceField(