Manage document generation for treatment and file treatment administrative act

2 files changed, 20 insertions, 7 deletions

diff --git a/archaeological_operations/forms.py b/archaeological_operations/forms.py
index b3f203e53..edd25da60 100644
--- a/archaeological_operations/forms.py
+++ b/archaeological_operations/forms.py
@@ -1403,8 +1403,8 @@ class DocumentGenerationAdminActForm(forms.Form):
 
     def __init__(self, *args, **kwargs):
         self.document_type = 'O'
-        if 'operation' in kwargs:
-            self.document_type = 'O' if kwargs.pop('operation') else 'F'
+        if 'document_type' in kwargs:
+            self.document_type = kwargs.pop('document_type')
         self.obj = None
         if 'obj' in kwargs:
             self.obj = kwargs.pop('obj')
diff --git a/archaeological_operations/views.py b/archaeological_operations/views.py
index 18d57e617..7a9c7329b 100644
--- a/archaeological_operations/views.py
+++ b/archaeological_operations/views.py
@@ -404,17 +404,30 @@ def generatedoc_administrativeactop(request, pk, template_pk=None):
     return HttpResponse(mimetype='text/plain')
 
 
-def administrativeactfile_document(request, operation=True):
+def administrativeactfile_document(
+        request, file=False, treatment=False, treatment_file=False):
     search_form = AdministrativeActOpeFormSelection
-    if not operation:
+    document_type = 'O'
+    if file:
         from archaeological_files.forms import \
             AdministrativeActFileFormSelection
         search_form = AdministrativeActFileFormSelection
+        document_type = 'F'
+    elif treatment:
+        from archaeological_finds.forms import \
+            AdministrativeActTreatmentFormSelection
+        search_form = AdministrativeActTreatmentFormSelection
+        document_type = 'T'
+    elif treatment_file:
+        from archaeological_finds.forms import \
+            AdministrativeActTreatmentFileFormSelection
+        search_form = AdministrativeActTreatmentFileFormSelection
+        document_type = 'TF'
     dct = {}
     if request.POST:
         dct['search_form'] = search_form(request.POST)
         dct['template_form'] = DocumentGenerationAdminActForm(
-            operation=operation)
+            document_type=document_type)
         c_object = None
         try:
             if dct['search_form'].is_valid():
@@ -425,7 +438,7 @@ def administrativeactfile_document(request, operation=True):
             pass
         if c_object:
             dct['template_form'] = DocumentGenerationAdminActForm(
-                request.POST, operation=operation, obj=c_object)
+                request.POST, document_type=document_type, obj=c_object)
             if dct['template_form'].is_valid():
                 return generatedoc_administrativeactop(
                     request,
@@ -434,7 +447,7 @@ def administrativeactfile_document(request, operation=True):
     else:
         dct['search_form'] = search_form()
         dct['template_form'] = DocumentGenerationAdminActForm(
-            operation=operation)
+            document_type=document_type)
     return render_to_response('ishtar/administrativeact_document.html', dct,
                               context_instance=RequestContext(request))