Containers / Warehouse: actions, sheets and fixes.

1 files changed, 11 insertions, 3 deletions

diff --git a/archaeological_warehouse/wizards.py b/archaeological_warehouse/wizards.py
index 407a58ad6..cfe5be4d4 100644
--- a/archaeological_warehouse/wizards.py
+++ b/archaeological_warehouse/wizards.py
@@ -17,13 +17,12 @@
 
 # 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
 
 from ishtar_common.wizards import Wizard
 from archaeological_finds.wizards import TreatmentWizard
-from archaeological_finds.models import Treatment
+from archaeological_finds.models import Treatment, TreatmentType
 
 import models
 
@@ -36,11 +35,20 @@ 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