Wizard: display selected on final wizard panel

1 files changed, 14 insertions, 6 deletions

diff --git a/archaeological_finds/wizards.py b/archaeological_finds/wizards.py
index e90fa745c..3314759d1 100644
--- a/archaeological_finds/wizards.py
+++ b/archaeological_finds/wizards.py
@@ -119,14 +119,17 @@ class TreatmentWizard(Wizard):
         kwargs['user'] = self.request.user
         return kwargs
 
-    def get_current_find(self):
+    def get_current_finds(self):
         step = self.steps.current
         if not step:
             return
         find_form_key = 'selecfind-treatment_creation'
-        find_id = self.session_get_value(find_form_key, "resulting_pk")
+        find_ids = self.session_get_value(find_form_key, "resulting_pk")
         try:
-            return models.Find.objects.get(pk=int(find_id))
+            return [
+                models.Find.objects.get(pk=int(find_id.strip()))
+                for find_id in find_ids.split(u',')
+            ]
         except(TypeError, ValueError, ObjectDoesNotExist):
             pass
 
@@ -134,12 +137,17 @@ class TreatmentWizard(Wizard):
         initial = super(TreatmentWizard, self).get_form_initial(step)
         if step != 'basetreatment-treatment_creation':
             return initial
-        find = self.get_current_find()
-        if not find or not find.container:
+        finds = self.get_current_finds()
+        if not finds:
+            return initial
+        locations = [find.container.location.pk for find in finds
+                     if find.container]
+        # no location or multiple locations
+        if not locations or len(set(locations)) != 1:
             return initial
         if not initial:
             initial = {}
-        initial['location'] = find.container.location.pk
+        initial['location'] = locations[0]
         return initial
 
     def get_extra_model(self, dct, form_list):