Treatment n<->1: manage saving

1 files changed, 134 insertions, 13 deletions

diff --git a/archaeological_finds/wizards.py b/archaeological_finds/wizards.py
index 8a7a3a513..3043d22aa 100644
--- a/archaeological_finds/wizards.py
+++ b/archaeological_finds/wizards.py
@@ -106,24 +106,16 @@ class TreatmentSearch(SearchWizard):
     model = models.Treatment
 
 
-class TreatmentWizard(Wizard):
+class TreatmentBase(Wizard):
     model = models.Treatment
     wizard_done_window = reverse_lazy('show-treatment')
-    basket_step = 'basetreatment-treatment_creation'
-    saved_args = {"items": []}
-
-    def get_form_kwargs(self, step, **kwargs):
-        kwargs = super(TreatmentWizard, self).get_form_kwargs(step, **kwargs)
-        if self.basket_step not in step:
-            return kwargs
-        kwargs['user'] = self.request.user
-        return kwargs
+    base_url = ""
 
     def get_current_finds(self):
         step = self.steps.current
         if not step:
             return
-        find_form_key = 'selecfind-treatment_creation'
+        find_form_key = 'selecfind-' + self.base_url
         find_ids = self.session_get_value(find_form_key, "resulting_pk")
         try:
             return [
@@ -134,8 +126,9 @@ class TreatmentWizard(Wizard):
             pass
 
     def get_form_initial(self, step, data=None):
-        initial = super(TreatmentWizard, self).get_form_initial(step)
-        if step != 'basetreatment-treatment_creation':
+        initial = super(TreatmentBase, self).get_form_initial(step)
+        base_step = 'basetreatment-' + self.base_url
+        if step != base_step:
             return initial
         finds = self.get_current_finds()
         if not finds:
@@ -150,6 +143,19 @@ class TreatmentWizard(Wizard):
         initial['location'] = locations[0]
         return initial
 
+
+class TreatmentWizard(TreatmentBase):
+    basket_step = 'basetreatment-treatment_creation'
+    saved_args = {"items": []}
+    base_url = 'treatment_creation'
+
+    def get_form_kwargs(self, step, **kwargs):
+        kwargs = super(TreatmentWizard, self).get_form_kwargs(step, **kwargs)
+        if self.basket_step not in step:
+            return kwargs
+        kwargs['user'] = self.request.user
+        return kwargs
+
     def get_extra_model(self, dct, form_list):
         """
         Get items concerned by the treatment
@@ -181,6 +187,121 @@ class TreatmentModificationWizard(TreatmentWizard):
     modification = True
 
 
+class TreatmentN1Wizard(TreatmentBase):
+    saved_args = {"upstream_items": [], "resulting_find": None}
+    base_url = 'treatment_creation_n1'
+
+    def _update_simple_initial_from_finds(self, initial, find, k):
+        r_k = "resulting_" + k
+        if getattr(find, k) and r_k in initial:
+            # pop k when value is inconsistent between finds
+            if initial[r_k] and initial[r_k] != getattr(find, k).pk:
+                initial.pop(r_k)
+            else:
+                initial[r_k] = getattr(find, k).pk
+        return initial
+
+    def _update_multi_initial_from_finds(self, initial, find, k):
+        r_k = "resulting_" + k
+        for value in getattr(find, k + 's').all():
+            if value.pk not in initial[r_k]:
+                initial[r_k].append(value.pk)
+        return initial
+
+    def _update_num_initial_from_finds(self, initial, find, k):
+        r_k = "resulting_" + k
+        if not getattr(find, k):
+            return initial
+        if initial[r_k] is None:
+            initial[r_k] = 0
+        initial[r_k] += getattr(find, k)
+        return initial
+
+    def _update_char_initial_from_finds(self, initial, find, k, sep=' ; '):
+        r_k = "resulting_" + k
+        if not getattr(find, k):
+            return initial
+        if not initial[r_k]:
+            initial[r_k] = getattr(find, k)
+        else:
+            initial[r_k] += sep + getattr(find, k)
+        return initial
+
+    def get_form_initial(self, step, data=None):
+        initial = super(TreatmentN1Wizard, self).get_form_initial(step)
+        if step != 'resultingfind-treatment_creation_n1':
+            return initial
+        finds = self.get_current_finds()
+        if not finds:
+            return initial
+        simple_key = ['material_type_quality']
+        multi_key = ['material_type', 'object_type', 'communicabilitie']
+        numeric_key = ['find_number', 'min_number_of_individuals']
+        desc_key = ['decoration', 'inscription', 'comment', 'dating_comment']
+        char_key = ['manufacturing_place']
+
+        for k in simple_key + numeric_key + desc_key + char_key:
+            initial["resulting_" + k] = None
+        for k in multi_key:
+            initial["resulting_" + k] = []
+
+        for find in finds:
+            for k in simple_key:
+                initial = self._update_simple_initial_from_finds(
+                    initial, find, k)
+            for k in multi_key:
+                initial = self._update_multi_initial_from_finds(
+                    initial, find, k)
+            for k in numeric_key:
+                initial = self._update_num_initial_from_finds(
+                    initial, find, k)
+            for k in char_key:
+                initial = self._update_char_initial_from_finds(
+                    initial, find, k, sep=' ; ')
+            for k in desc_key:
+                initial = self._update_char_initial_from_finds(
+                    initial, find, k, sep='\n')
+
+        for k in initial.keys():
+            if initial[k] is None:
+                initial.pop(k)
+        return initial
+
+    def get_extra_model(self, dct, form_list):
+        """
+        Get items concerned by the treatment
+        """
+        dct = super(TreatmentN1Wizard, self).get_extra_model(dct, form_list)
+        if 'resulting_pk' not in dct:
+            return dct
+
+        dct['upstream_items'] = []
+        # manage upstream items
+        pks = dct.pop('resulting_pk')
+        if hasattr(pks, 'split'):
+            pks = pks.split(u',')  # unicode or string
+        for pk in pks:
+            try:
+                find = models.Find.objects.get(pk=pk)
+                dct['upstream_items'].append(find)
+            except models.Find.DoesNotExist:
+                raise PermissionDenied
+
+        for find in dct['upstream_items']:
+            if 'own' in self.current_right \
+                    and not find.is_own(dct['history_modifier']):
+                raise PermissionDenied
+
+        # extract data of the new find
+        dct['resulting_find'] = {}
+        for k in dct.keys():
+            if k.startswith('resulting_') and k != "resulting_find":
+                dct['resulting_find'][
+                    k[len('resulting_'):]
+                ] = dct.pop(k)
+        return dct
+
+
 class TreatmentDeletionWizard(DeletionWizard):
     model = models.Treatment
     fields = ['label', 'other_reference', 'year', 'index',