Container: Localisation form become modification form

1 files changed, 12 insertions, 11 deletions

diff --git a/archaeological_warehouse/wizards.py b/archaeological_warehouse/wizards.py
index bdb3959dc..4bdc150cc 100644
--- a/archaeological_warehouse/wizards.py
+++ b/archaeological_warehouse/wizards.py
@@ -66,18 +66,20 @@ class WarehouseModificationWizard(Wizard):
     wizard_done_window = reverse_lazy('show-warehouse')
 
 
-class LocalisationWizard(Wizard):
+class ContainerWizard(Wizard):
+    model = models.Container
     wizard_templates = {
         'localisation-container_localisation':
             'ishtar/wizard/wizard_containerlocalisation.html',
     }
+    ignore_init_steps = ['localisation']
 
     def get_form_kwargs(self, step=None):
-        kwargs = super(LocalisationWizard, self).get_form_kwargs(step)
-        if step == 'localisation-container_localisation':
+        kwargs = super(ContainerWizard, self).get_form_kwargs(step)
+        if step == 'localisation-' + self.url_name:
             container = None
             container_pk = self.session_get_value(
-                'seleccontainer-container_localisation', 'container')
+                'selec-' + self.url_name, 'pk')
             q = models.Container.objects.filter(pk=container_pk)
             if q.count():
                 container = q.all()[0]
@@ -88,17 +90,13 @@ class LocalisationWizard(Wizard):
         """
         Save the localisation
         """
+        super(ContainerWizard, self).done(form_list)
         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
+            container = self.get_current_object()
+            if idx == 2:  # localisation form
                 for div_name in form.cleaned_data:
                     try:
                         division = models.WarehouseDivisionLink.objects.get(
@@ -117,3 +115,6 @@ class LocalisationWizard(Wizard):
         return render_to_response(
             self.wizard_done_template, {},
             context_instance=RequestContext(self.request))
+
+class ContainerModificationWizard(ContainerWizard):
+    modification = True