1 files changed, 130 insertions, 3 deletions
diff --git a/archaeological_finds/wizards.py b/archaeological_finds/wizards.py
index f897969c4..c9d492329 100644
--- a/archaeological_finds/wizards.py
+++ b/archaeological_finds/wizards.py
@@ -17,11 +17,14 @@
 
 # See the file COPYING for details.
 
-from django.core.exceptions import ObjectDoesNotExist
+from django.core.exceptions import ObjectDoesNotExist, PermissionDenied
 from django.utils.translation import ugettext_lazy as _
 
 from ishtar_common.forms import reverse_lazy
 from ishtar_common.wizards import Wizard, DeletionWizard, SourceWizard
+from archaeological_operations.wizards import OperationAdministrativeActWizard
+
+from archaeological_operations.models import AdministrativeAct
 import models
 
 
@@ -84,15 +87,125 @@ class FindDeletionWizard(DeletionWizard):
 
 class TreatmentWizard(Wizard):
     model = models.Treatment
+    wizard_done_window = reverse_lazy('show-treatment')
     basket_step = 'basetreatment-treatment_creation'
+    saved_args = {"items": []}
 
-    def get_form_kwargs(self, step):
-        kwargs = super(TreatmentWizard, self).get_form_kwargs(step)
+    def get_form_kwargs(self, step, **kwargs):
+        kwargs = super(TreatmentWizard, self).get_form_kwargs(step, **kwargs)
         if self.basket_step not in step:
             return kwargs
         kwargs['user'] = self.request.user
         return kwargs
 
+    def get_extra_model(self, dct, form_list):
+        """
+        Remove basket ID to the result dict
+        """
+        dct = super(TreatmentWizard, self).get_extra_model(dct, form_list)
+        if 'resulting_pk' in dct:
+            try:
+                find = models.Find.objects.get(pk=dct.pop('resulting_pk'))
+                if 'own' in self.current_right \
+                        and not find.is_own(dct['history_modifier']):
+                    raise PermissionDenied
+                dct['items'] = [find]
+            except models.Find.DoesNotExist:
+                raise PermissionDenied
+        if 'basket' in dct:
+            basket = dct.pop('basket')
+            if basket.user.pk != dct['history_modifier'].pk:
+                raise PermissionDenied
+            dct['items'] = list(basket.items.all())
+        return dct
+
+
+class TreatmentModificationWizard(TreatmentWizard):
+    modification = True
+
+
+class TreatmentDeletionWizard(DeletionWizard):
+    model = models.Treatment
+    fields = ['label', 'other_reference', 'year', 'index',
+              'treatment_types', 'location', 'person', 'organization',
+              'external_id', 'comment', 'description',
+              'goal', 'start_date', 'end_date', 'container']
+
+
+class TreatmentAdministrativeActWizard(OperationAdministrativeActWizard):
+    model = models.Treatment
+    current_obj_slug = 'administrativeacttreatment'
+    ref_object_key = 'treatment'
+
+    def get_reminder(self):
+        return
+
+
+class TreatmentEditAdministrativeActWizard(TreatmentAdministrativeActWizard):
+    model = AdministrativeAct
+    edit = True
+
+    def get_associated_item(self, dct):
+        return self.get_current_object().treatment
+
+
+class TreatmentFileWizard(Wizard):
+    model = models.TreatmentFile
+    wizard_done_window = reverse_lazy('show-treatmentfile')
+
+
+class TreatmentFileModificationWizard(TreatmentFileWizard):
+    modification = True
+
+
+class TreatmentFileDeletionWizard(DeletionWizard):
+    model = models.TreatmentFile
+    fields = ['name', 'internal_reference', 'external_id', 'year',
+              'index', 'type', 'in_charge', 'reception_date',
+              'creation_date', 'end_date', 'comment']
+
+
+class TreatmentFileAdministrativeActWizard(
+        OperationAdministrativeActWizard):
+    model = models.TreatmentFile
+    current_obj_slug = 'administrativeacttreatmentfile'
+    ref_object_key = 'treatment_file'
+
+    def get_reminder(self):
+        form_key = 'selec-' + self.url_name
+        if self.url_name.endswith('_administrativeactop'):
+            # modification and deletion are suffixed with '_modification'
+            # and '_deletion' so it is creation
+            pk = self.session_get_value(form_key, "pk")
+            try:
+                return (
+                    (_(u"Treatment file"),
+                     unicode(models.Operation.objects.get(pk=pk))),
+                )
+            except models.TreatmentFile.DoesNotExist:
+                return
+        else:
+            admin_id = self.session_get_value(form_key, "pk")
+            try:
+                admin = AdministrativeAct.objects.get(pk=admin_id)
+                if not admin.operation:
+                    return
+                return (
+                    (_(u"Operation"), unicode(admin.operation)),
+                )
+            except AdministrativeAct.DoesNotExist:
+                return
+        return
+
+
+class TreatmentFileEditAdministrativeActWizard(
+        TreatmentFileAdministrativeActWizard):
+    model = AdministrativeAct
+    edit = True
+
+    def get_associated_item(self, dct):
+        return self.get_current_object().treatment_file
+
 
 class FindSourceWizard(SourceWizard):
     wizard_done_window = reverse_lazy('show-findsource')
@@ -106,3 +219,17 @@ class FindSourceDeletionWizard(DeletionWizard):
 
 class TreatmentSourceWizard(SourceWizard):
     model = models.TreatmentSource
+
+
+class TreatmentSourceDeletionWizard(DeletionWizard):
+    model = models.TreatmentSource
+    fields = ['treatment', 'title', 'source_type', 'authors']
+
+
+class TreatmentFileSourceWizard(SourceWizard):
+    model = models.TreatmentFileSource
+
+
+class TreatmentFileSourceDeletionWizard(DeletionWizard):
+    model = models.TreatmentFileSource
+    fields = ['treatment_file', 'title', 'source_type', 'authors']