Container creation form (refs #3392)

1 files changed, 8 insertions, 6 deletions

diff --git a/archaeological_warehouse/wizards.py b/archaeological_warehouse/wizards.py
index 4bdc150cc..0cc6915f8 100644
--- a/archaeological_warehouse/wizards.py
+++ b/archaeological_warehouse/wizards.py
@@ -77,13 +77,16 @@ class ContainerWizard(Wizard):
     def get_form_kwargs(self, step=None):
         kwargs = super(ContainerWizard, self).get_form_kwargs(step)
         if step == 'localisation-' + self.url_name:
-            container = None
             container_pk = self.session_get_value(
                 'selec-' + self.url_name, 'pk')
             q = models.Container.objects.filter(pk=container_pk)
             if q.count():
-                container = q.all()[0]
-            kwargs['container'] = container
+                kwargs['container'] = q.all()[0]
+            warehouse_pk = self.session_get_value(
+                'container-' + self.url_name, 'location')
+            q = models.Warehouse.objects.filter(pk=warehouse_pk)
+            if q.count():
+                kwargs['warehouse'] = q.all()[0]
         return kwargs
 
     def done(self, form_list, **kwargs):
@@ -91,12 +94,11 @@ class ContainerWizard(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)
-            container = self.get_current_object()
-            if idx == 2:  # localisation form
+            container = self.get_current_object() or self.current_object
+            if container and form.prefix == 'localisation-' + self.url_name:
                 for div_name in form.cleaned_data:
                     try:
                         division = models.WarehouseDivisionLink.objects.get(