QA find: add divide action

1 files changed, 44 insertions, 0 deletions

diff --git a/archaeological_finds/views.py b/archaeological_finds/views.py
index 9e1a22a2a..f0dfe28e7 100644
--- a/archaeological_finds/views.py
+++ b/archaeological_finds/views.py
@@ -591,12 +591,56 @@ def treatment_add(request, pks, treatment_file=None):
         kwargs={'step': 'basetreatment-treatment_creation'}))
 
 
+def divide_treatment_add(request, pks, treatment_file=None):
+    treatment_creation_1n_wizard(request)
+
+    wizards.Treatment1NWizard.session_set_value(
+        request, 'selecfind-treatment_creation_1n',
+        'resulting_pk', pks, reset=True)
+    if treatment_file:
+        wizards.Treatment1NWizard.session_set_value(
+            request, 'file-treatment_creation_1n', 'file', treatment_file.pk)
+    else:
+        wizards.Treatment1NWizard.session_set_value(
+            request, 'file-treatment_creation_1n', '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 and treatment_file.type.treatment_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.Treatment1NWizard.session_set_value(
+                request, 'basetreatment-treatment_creation_1n', k, dct[k])
+    return redirect(reverse(
+        'treatment_creation_1n',
+        kwargs={'step': 'basetreatment-treatment_creation_1n'}))
+
+
 def find_treatment_add(request, pk, current_right=None):
     if not models.Find.objects.filter(pk=pk).count():
         raise Http404()
     return treatment_add(request, str(pk))
 
 
+def find_divide_treatment_add(request, pk, current_right=None):
+    if not models.Find.objects.filter(pk=pk).count():
+        raise Http404()
+    return divide_treatment_add(request, str(pk))
+
+
 def findbasket_treatment_add(request, pk, current_right=None):
     try:
         basket = models.FindBasket.objects.get(pk=pk)