1 files changed, 54 insertions, 0 deletions
diff --git a/archaeological_warehouse/wizards.py b/archaeological_warehouse/wizards.py
index f5aee5c1e..407a58ad6 100644
--- a/archaeological_warehouse/wizards.py
+++ b/archaeological_warehouse/wizards.py
@@ -17,6 +17,7 @@
 
 # See the file COPYING for details.
 
+from django.contrib.formtools.wizard.views import NamedUrlWizardView
 from django.shortcuts import render_to_response
 from django.template import RequestContext
 
@@ -52,3 +53,56 @@ class WarehouseWizard(Wizard):
 class WarehouseModificationWizard(Wizard):
     model = models.Warehouse
     modification = True
+
+
+class LocalisationWizard(Wizard):
+    wizard_templates = {
+        'localisation-container_localisation':
+            'ishtar/wizard/wizard_containerlocalisation.html',
+    }
+
+    def get_form_kwargs(self, step=None):
+        kwargs = super(LocalisationWizard, self).get_form_kwargs(step)
+        if step == 'localisation-container_localisation':
+            container = None
+            container_pk = self.session_get_value(
+                'seleccontainer-container_localisation', 'container')
+            q = models.Container.objects.filter(pk=container_pk)
+            if q.count():
+                container = q.all()[0]
+            kwargs['container'] = container
+        return kwargs
+
+    def done(self, form_list, **kwargs):
+        """
+        Save the localisation
+        """
+        container = None
+        for idx, form in enumerate(form_list):
+            if not form.is_valid():
+                return self.render(form)
+            if not idx:  # container form
+                try:
+                    container = models.Container.objects.get(
+                        pk=form.cleaned_data['container'])
+                except models.Container.DoesNotExist:
+                    return self.render(form)
+            elif idx == 1:  # localisation form
+                for div_name in form.cleaned_data:
+                    try:
+                        division = models.WarehouseDivisionLink.objects.get(
+                            pk=div_name.split('_')[-1],
+                            warehouse=container.location
+                        )  # check the localisation match with the container
+                    except models.WarehouseDivisionLink.DoesNotExist:
+                        return self.render(form)
+                    localisation, created = \
+                        models.ContainerLocalisation.objects.get_or_create(
+                            container=container,
+                            division=division
+                        )
+                    localisation.reference = form.cleaned_data[div_name]
+                    localisation.save()
+        return render_to_response(
+            self.wizard_done_template, {},
+            context_instance=RequestContext(self.request))