Warehouse: work on packaging

1 files changed, 12 insertions, 5 deletions

diff --git a/archaeological_warehouse/wizards.py b/archaeological_warehouse/wizards.py
index 833c1c697..5f2b4346d 100644
--- a/archaeological_warehouse/wizards.py
+++ b/archaeological_warehouse/wizards.py
@@ -1,6 +1,6 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
-# Copyright (C) 2012  Étienne Loks  <etienne.loks_AT_peacefrogsDOTnet>
+# Copyright (C) 2012-2016  Étienne Loks  <etienne.loks_AT_peacefrogsDOTnet>
 
 # This program is free software: you can redistribute it and/or modify
 # it under the terms of the GNU Affero General Public License as
@@ -25,17 +25,24 @@ from archaeological_finds.models import Treatment
 
 
 class PackagingWizard(TreatmentWizard):
+
+    def get_form_kwargs(self, step):
+        kwargs = super(PackagingWizard, self).get_form_kwargs(step)
+        if 'base-packaging' not in step:
+            return kwargs
+        kwargs['user'] = self.request.user
+        return kwargs
+
+
     def save_model(self, dct, m2m, whole_associated_models, form_list,
                    return_object):
         dct = self.get_extra_model(dct, form_list)
         obj = self.get_current_saved_object()
         dct['location'] = dct['container'].location
-        items = dct.pop('finds')
+        basket = dct.pop('basket')
         treatment = Treatment(**dct)
         treatment.save()
-        if not hasattr(items, '__iter__'):
-            items = [items]
-        for item in items:
+        for item in basket.items.all():
             new = item.duplicate(self.request.user)
             item.downstream_treatment = treatment
             item.save()