Manage treatment with no creation new of item

1 files changed, 26 insertions, 12 deletions

diff --git a/archaeological_finds/wizards.py b/archaeological_finds/wizards.py
index bd16a55de..29ca50056 100644
--- a/archaeological_finds/wizards.py
+++ b/archaeological_finds/wizards.py
@@ -86,8 +86,8 @@ class FindWizard(Wizard):
             (_(u"Context record"), unicode(current_cr)))
         return context
 
-    def get_extra_model(self, dct, form_list):
-        dct = super(FindWizard, self).get_extra_model(dct, form_list)
+    def get_extra_model(self, dct, m2m, form_list):
+        dct = super(FindWizard, self).get_extra_model(dct, m2m, form_list)
         dct['order'] = 1
         if 'pk' in dct and type(dct['pk']) == ContextRecord:
             dct['base_finds__context_record'] = dct.pop('pk')
@@ -125,6 +125,7 @@ class TreatmentBase(Wizard):
     model = models.Treatment
     wizard_done_window = reverse_lazy('show-treatment')
     base_url = ""
+    saved_args = {"treatment_type_list": []}
 
     def get_current_finds(self):
         step = self.steps.current
@@ -158,10 +159,20 @@ class TreatmentBase(Wizard):
         initial['location'] = locations[0]
         return initial
 
+    def get_extra_model(self, dct, m2m, form_list):
+        dct = super(TreatmentBase, self).get_extra_model(dct, m2m, form_list)
+        dct['treatment_type_list'] = []
+        for k, v in m2m:
+            if k == 'treatment_type':
+                if type(v) not in (list, tuple):
+                    v = [v]
+                dct['treatment_type_list'] += v
+        return dct
+
 
 class TreatmentWizard(TreatmentBase):
     basket_step = 'basetreatment-treatment_creation'
-    saved_args = {"items": []}
+    saved_args = {"items": [], "treatment_type_list": []}
     base_url = 'treatment_creation'
 
     def get_form_kwargs(self, step, **kwargs):
@@ -171,11 +182,11 @@ class TreatmentWizard(TreatmentBase):
         kwargs['user'] = self.request.user
         return kwargs
 
-    def get_extra_model(self, dct, form_list):
+    def get_extra_model(self, dct, m2m, form_list):
         """
         Get items concerned by the treatment
         """
-        dct = super(TreatmentWizard, self).get_extra_model(dct, form_list)
+        dct = super(TreatmentWizard, self).get_extra_model(dct, m2m, form_list)
         if 'resulting_pk' in dct:
             dct['items'] = []
             pks = dct.pop('resulting_pk').split(u',')
@@ -203,7 +214,8 @@ class TreatmentModificationWizard(TreatmentWizard):
 
 
 class TreatmentN1Wizard(TreatmentBase):
-    saved_args = {"upstream_items": [], "resulting_find": None}
+    saved_args = {"upstream_items": [], "resulting_find": None,
+                  "treatment_type_list": []}
     base_url = 'treatment_creation_n1'
 
     def _update_simple_initial_from_finds(self, initial, find, k):
@@ -282,11 +294,12 @@ class TreatmentN1Wizard(TreatmentBase):
                 initial.pop(k)
         return initial
 
-    def get_extra_model(self, dct, form_list):
+    def get_extra_model(self, dct, m2m, form_list):
         """
         Get items concerned by the treatment
         """
-        dct = super(TreatmentN1Wizard, self).get_extra_model(dct, form_list)
+        dct = super(TreatmentN1Wizard, self).get_extra_model(
+            dct, m2m, form_list)
         if 'resulting_pk' not in dct:
             return dct
 
@@ -318,7 +331,8 @@ class TreatmentN1Wizard(TreatmentBase):
 
 
 class Treatment1NWizard(TreatmentBase):
-    saved_args = {"upstream_item": None, "resulting_finds": None}
+    saved_args = {"upstream_item": None, "resulting_finds": None,
+                  "treatment_type_list": []}
     base_url = 'treatment_creation_1n'
 
     def get_form_kwargs(self, step, **kwargs):
@@ -340,11 +354,12 @@ class Treatment1NWizard(TreatmentBase):
         initial['resultings_label'] = lbl + u"-"
         return initial
 
-    def get_extra_model(self, dct, form_list):
+    def get_extra_model(self, dct, m2m, form_list):
         """
         Get items concerned by the treatment
         """
-        dct = super(Treatment1NWizard, self).get_extra_model(dct, form_list)
+        dct = super(Treatment1NWizard, self).get_extra_model(
+            dct, m2m, form_list)
         if 'resulting_pk' not in dct:
             return dct
 
@@ -373,7 +388,6 @@ class Treatment1NWizard(TreatmentBase):
                       u"resulting items. This search have been pinned.")
                     ).format(dct["resulting_finds"]["basket_name"])
         )
-
         self.request.session["pin-search-find"] = u'{}="{}"'.format(
             unicode(pgettext("key for text search", u"basket")),
             dct["resulting_finds"]["basket_name"])