Reorganize treatment actions - add a direct to treatements

1 files changed, 23 insertions, 0 deletions

diff --git a/archaeological_finds/wizards.py b/archaeological_finds/wizards.py
index 0e0fe80e1..e90fa745c 100644
--- a/archaeological_finds/wizards.py
+++ b/archaeological_finds/wizards.py
@@ -119,6 +119,29 @@ class TreatmentWizard(Wizard):
         kwargs['user'] = self.request.user
         return kwargs
 
+    def get_current_find(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")
+        try:
+            return models.Find.objects.get(pk=int(find_id))
+        except(TypeError, ValueError, ObjectDoesNotExist):
+            pass
+
+    def get_form_initial(self, step, data=None):
+        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:
+            return initial
+        if not initial:
+            initial = {}
+        initial['location'] = find.container.location.pk
+        return initial
+
     def get_extra_model(self, dct, form_list):
         """
         Get items concerned by the treatment