diff options
Diffstat (limited to 'archaeological_finds/wizards.py')
-rw-r--r-- | archaeological_finds/wizards.py | 20 |
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): |