1 files changed, 93 insertions, 2 deletions
diff --git a/archaeological_warehouse/wizards.py b/archaeological_warehouse/wizards.py
index 62b441e00..817ed75a5 100644
--- a/archaeological_warehouse/wizards.py
+++ b/archaeological_warehouse/wizards.py
@@ -20,8 +20,12 @@
 from django.shortcuts import render_to_response
 from django.template import RequestContext
 
+from ishtar_common.forms import reverse_lazy
+from ishtar_common.wizards import Wizard, DeletionWizard
 from archaeological_finds.wizards import TreatmentWizard
-from archaeological_finds.models import Treatment
+
+from archaeological_finds.models import Treatment, TreatmentType
+import models
 
 
 class PackagingWizard(TreatmentWizard):
@@ -32,11 +36,98 @@ class PackagingWizard(TreatmentWizard):
         dct = self.get_extra_model(dct, form_list)
         obj = self.get_current_saved_object()
         dct['location'] = dct['container'].location
-        items = dct.pop('basket')
+        items = None
+        if 'items' in dct:
+            items = dct.pop('items')
+        if 'basket' in dct:
+            if not items:
+                items = dct.pop('basket')
+            else:
+                dct.pop('basket')
         treatment = Treatment(**dct)
         extra_args_for_new = {"container": dct['container']}
         treatment.save(items=items, user=self.request.user,
                        extra_args_for_new=extra_args_for_new)
+        packaging = TreatmentType.objects.get(txt_idx='packaging')
+        treatment.treatment_types.add(packaging)
         res = render_to_response('ishtar/wizard/wizard_done.html', {},
                                  context_instance=RequestContext(self.request))
         return return_object and (obj, res) or res
+
+
+class WarehouseWizard(Wizard):
+    model = models.Warehouse
+    wizard_done_window = reverse_lazy('show-warehouse')
+
+
+class WarehouseModificationWizard(Wizard):
+    model = models.Warehouse
+    modification = True
+    wizard_done_window = reverse_lazy('show-warehouse')
+
+
+class WarehouseDeletionWizard(DeletionWizard):
+    model = models.Warehouse
+
+
+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(ContainerWizard, self).get_form_kwargs(step)
+        if step == 'localisation-' + self.url_name:
+            container_pk = self.session_get_value(
+                'selec-' + self.url_name, 'pk')
+            q = models.Container.objects.filter(pk=container_pk)
+            if q.count():
+                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):
+        """
+        Save the localisation
+        """
+        super(ContainerWizard, self).done(form_list)
+        for idx, form in enumerate(form_list):
+            if not form.is_valid():
+                return self.render(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(
+                            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))
+
+
+class ContainerModificationWizard(ContainerWizard):
+    modification = True
+
+
+class ContainerDeletionWizard(DeletionWizard):
+    model = models.Container
+    fields = ['container_type', 'reference', 'comment', 'location', 'index',
+              'cached_label']