Operation: easily add all parcel related to an archaeological file (refs #1444)

1 files changed, 45 insertions, 12 deletions

diff --git a/archaeological_operations/wizards.py b/archaeological_operations/wizards.py
index 91f4be3d4..48f7aeb65 100644
--- a/archaeological_operations/wizards.py
+++ b/archaeological_operations/wizards.py
@@ -51,11 +51,16 @@ class OperationWizard(Wizard):
         """
         context = super(OperationWizard, self).get_context_data(form,
                                                                 **kwargs)
-        #step = self.determine_step(request, storage)
         step = self.steps.current
-        if not step.startswith('towns-'):
-            return context
-        context['TOWNS'] = self.get_towns()
+        if step.startswith('towns-'):
+            context['TOWNS'] = self.get_towns()
+        elif step.startswith('parcels-'):
+            # if a file is acciated to the operation add the button "Add all"
+            general_form_key = 'general-' + self.url_name
+            file_id = self.session_get_value(general_form_key,
+                                             "associated_file")
+            if file_id:
+                context['add_all'] = True
         return context
 
     def get_towns(self):
@@ -76,6 +81,16 @@ class OperationWizard(Wizard):
         else:
             return -1
 
+    def get_available_parcels(self, file_id):
+        parcels = []
+        try:
+            for parcel in models.File.objects.get(pk=int(file_id)
+                                                ).parcels.all():
+                parcels.append((parcel.pk, parcel.short_label))
+        except (ValueError, ObjectDoesNotExist):
+            pass
+        return sorted(parcels, key=lambda x:x[1])
+
     def get_form(self, step=None, data=None, files=None):
         """
         Manage specifics fields
@@ -96,14 +111,7 @@ class OperationWizard(Wizard):
             file_id = self.session_get_value(general_form_key,
                                              "associated_file")
             if file_id:
-                parcels = []
-                try:
-                    for parcel in models.File.objects.get(pk=int(file_id)
-                                                        ).parcels.all():
-                        parcels.append((parcel.pk, parcel.short_label))
-                except (ValueError, ObjectDoesNotExist):
-                    pass
-                data['PARCELS'] = sorted(parcels, key=lambda x:x[1])
+                data['PARCELS'] = self.get_available_parcels(file_id)
             else:
                 town_form_key = step.startswith('parcelsgeneral') \
                                 and 'townsgeneral-' or 'towns-'
@@ -165,6 +173,31 @@ class OperationWizard(Wizard):
             #self.form_initialized = True
         return initial'''
 
+    def post(self, *args, **kwargs):
+        request = self.request
+        post_data = request.POST.copy()
+
+        # add all parcel from available in the archaelogical file
+        if not post_data.get('add_all_parcels'):
+            return super(Wizard, self).post(*args, **kwargs)
+        general_form_key = 'general-' + self.url_name
+        file_id = self.session_get_value(general_form_key,
+                                         "associated_file")
+        if not file_id:
+            return super(Wizard, self).post(*args, **kwargs)
+        parcel_form_key = "parcels-" + self.url_name
+        idx = -1
+        # remove non relevant deleted keys
+        for k in post_data.keys():
+            if k.startswith(parcel_form_key) and k.endswith('-DELETE'):
+                post_data.pop(k)
+        for idx, parcel in enumerate(self.get_available_parcels(file_id)):
+            parcel_pk, parcel_name = parcel
+            post_data["%s-%d-parcel" % (parcel_form_key, idx)] = parcel_pk
+        post_data[parcel_form_key+'-TOTAL_FORMS'] = idx + 1
+        request.POST = post_data
+        return super(Wizard, self).post(*args, **kwargs)
+
 class OperationModificationWizard(OperationWizard):
     modification = True