Manage 1<->n treatments wizard

1 files changed, 66 insertions, 1 deletions

diff --git a/archaeological_finds/wizards.py b/archaeological_finds/wizards.py
index 3c25258a6..bd16a55de 100644
--- a/archaeological_finds/wizards.py
+++ b/archaeological_finds/wizards.py
@@ -17,8 +17,9 @@
 
 # See the file COPYING for details.
 
+from django.contrib import messages
 from django.core.exceptions import ObjectDoesNotExist, PermissionDenied
-from django.utils.translation import ugettext_lazy as _
+from django.utils.translation import ugettext_lazy as _, pgettext
 
 from ishtar_common.forms import reverse_lazy
 from ishtar_common.wizards import Wizard, DeletionWizard, SearchWizard
@@ -316,6 +317,70 @@ class TreatmentN1Wizard(TreatmentBase):
         return dct
 
 
+class Treatment1NWizard(TreatmentBase):
+    saved_args = {"upstream_item": None, "resulting_finds": None}
+    base_url = 'treatment_creation_1n'
+
+    def get_form_kwargs(self, step, **kwargs):
+        kwargs = super(Treatment1NWizard, self).get_form_kwargs(step, **kwargs)
+        if step != 'resultingfind-treatment_creation_1n':
+            return kwargs
+        kwargs['user'] = self.request.user
+        return kwargs
+
+    def get_form_initial(self, step, data=None):
+        initial = super(Treatment1NWizard, self).get_form_initial(step)
+        if step != 'resultingfinds-treatment_creation_1n':
+            return initial
+        finds = self.get_current_finds()
+        if not finds:
+            return initial
+        lbl = finds[0].label
+        initial['resultings_basket_name'] = unicode(_(u"Basket")) + u" - " + lbl
+        initial['resultings_label'] = lbl + u"-"
+        return initial
+
+    def get_extra_model(self, dct, form_list):
+        """
+        Get items concerned by the treatment
+        """
+        dct = super(Treatment1NWizard, self).get_extra_model(dct, form_list)
+        if 'resulting_pk' not in dct:
+            return dct
+
+        # manage upstream item
+        pk = dct.pop('resulting_pk')
+        try:
+            find = models.Find.objects.get(pk=pk)
+            dct['upstream_item'] = find
+        except models.Find.DoesNotExist:
+            raise PermissionDenied
+
+        if 'own' in self.current_right \
+                and not find.is_own(dct['history_modifier']):
+            raise PermissionDenied
+
+        # extract attributes to generate the new find
+        dct['resulting_finds'] = {}
+        for k in dct.keys():
+            if k.startswith('resultings_'):
+                dct['resulting_finds'][
+                    k[len('resultings_'):]
+                ] = dct.pop(k)
+        messages.add_message(
+            self.request, messages.INFO,
+            unicode(_(u"The new basket: \"{}\" have been created with the "
+                      u"resulting items. This search have been pinned.")
+                    ).format(dct["resulting_finds"]["basket_name"])
+        )
+
+        self.request.session["pin-search-find"] = u'{}="{}"'.format(
+            unicode(pgettext("key for text search", u"basket")),
+            dct["resulting_finds"]["basket_name"])
+        self.request.session['find'] = ''
+        return dct
+
+
 class TreatmentDeletionWizard(DeletionWizard):
     model = models.Treatment
     fields = ['label', 'other_reference', 'year', 'index',