Treatment file sheet: add treatment action

1 files changed, 41 insertions, 3 deletions

diff --git a/archaeological_finds/views.py b/archaeological_finds/views.py
index a04d7fe2f..a0964fa47 100644
--- a/archaeological_finds/views.py
+++ b/archaeological_finds/views.py
@@ -529,13 +529,37 @@ def treatment_modify(request, pk):
         kwargs={'step': 'basetreatment-treatment_modification'}))
 
 
-def treatment_add(request, pks):
+def treatment_add(request, pks, treatment_file=None):
     treatment_creation_wizard(request)
     wizards.TreatmentWizard.session_set_value(
         request, 'selecfind-treatment_creation',
         'resulting_pk', pks, reset=True)
-    wizards.TreatmentWizard.session_set_value(
-        request, 'file-treatment_creation', 'file', '')
+    if treatment_file:
+        wizards.TreatmentWizard.session_set_value(
+            request, 'file-treatment_creation', 'file', treatment_file.pk)
+    else:
+        wizards.TreatmentWizard.session_set_value(
+            request, 'file-treatment_creation', 'file', '')
+    if treatment_file:
+        in_charge = treatment_file.in_charge
+        if not in_charge:
+            in_charge = request.user.ishtaruser.person
+        dct = {
+            "treatment_type": treatment_file.type.treatment_type.pk
+            if treatment_file.type else "",
+            "year": treatment_file.year,
+            "person": in_charge.pk,
+        }
+        locas = list(
+            set([str(f.container.location.pk)
+                 for f in treatment_file.associated_basket.items.all()
+                 if f.container and f.container.location])
+        )
+        if len(locas) == 1:  # one and only one location for all finds
+            dct["location"] = locas[0]
+        for k in dct:
+            wizards.TreatmentWizard.session_set_value(
+                request, 'basetreatment-treatment_creation', k, dct[k])
     return redirect(reverse(
         'treatment_creation',
         kwargs={'step': 'basetreatment-treatment_creation'}))
@@ -565,6 +589,20 @@ def container_treatment_add(request, pk, current_right=None):
                          ",".join([str(f.pk) for f in basket.items.all()]))
 
 
+def treatmentfile_treatment_add(request, pk, current_right=None):
+    try:
+        tf = models.TreatmentFile.objects.get(pk=pk)
+    except models.TreatmentFile.DoesNotExist:
+        raise Http404()
+    if not tf.associated_basket:
+        raise Http404()
+    basket = tf.associated_basket
+    return treatment_add(
+        request, ",".join([str(f.pk) for f in basket.items.all()]),
+        treatment_file=tf
+    )
+
+
 treatment_deletion_wizard = wizards.TreatmentDeletionWizard.as_view([
     ('selec-treatment_deletion', forms.TreatmentFormSelection),
     ('final-treatment_deletion', forms.TreatmentDeletionForm)],