Simple treatment form. Treatment listing. (refs #3365)

1 files changed, 22 insertions, 3 deletions

diff --git a/archaeological_finds/wizards.py b/archaeological_finds/wizards.py
index f897969c4..e97c4518e 100644
--- a/archaeological_finds/wizards.py
+++ b/archaeological_finds/wizards.py
@@ -17,7 +17,7 @@
 
 # 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
@@ -85,14 +85,33 @@ class FindDeletionWizard(DeletionWizard):
 class TreatmentWizard(Wizard):
     model = models.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:
+            dct.pop('basket')
+        return dct
+
 
 class FindSourceWizard(SourceWizard):
     wizard_done_window = reverse_lazy('show-findsource')