Treatments: manage conditional forms

3 files changed, 97 insertions, 21 deletions

diff --git a/archaeological_finds/forms.py b/archaeological_finds/forms.py
index 9fdf2ddf2..997773761 100644
--- a/archaeological_finds/forms.py
+++ b/archaeological_finds/forms.py
@@ -283,6 +283,58 @@ class FindMultipleFormSelection(forms.Form):
         return self.cleaned_data
 
 
+def check_exist(form_name, key):
+    def func(self):
+        request = self.request
+        storage = self.storage
+        if storage.prefix not in request.session or \
+           'step_data' not in request.session[storage.prefix] or \
+           form_name not in request.session[storage.prefix]['step_data'] or\
+           form_name + '-' + key not in \
+           request.session[storage.prefix]['step_data'][form_name]:
+            return False
+        try:
+            val = request.session[storage.prefix][
+                'step_data'][form_name][form_name + '-' + key]
+            if val and len(val) > 0:
+                val = val[0]
+            return bool(val)
+        except ValueError:
+            return False
+    return func
+
+
+def check_not_exist(form_name, key):
+    def func(self):
+        return not check_exist(form_name, key)(self)
+    return func
+
+
+def check_type_field(form_name, key, model, field):
+    def func(self):
+        request = self.request
+        storage = self.storage
+        if storage.prefix not in request.session or \
+           'step_data' not in request.session[storage.prefix] or \
+           form_name not in request.session[storage.prefix]['step_data'] or\
+           form_name + '-' + key not in \
+           request.session[storage.prefix]['step_data'][form_name]:
+            return False
+        try:
+            val = model.objects.get(pk=request.session[storage.prefix][
+                'step_data'][form_name][form_name + '-' + key][0])
+            return bool(getattr(val, field))
+        except (ValueError, model.DoesNotExist):
+            return False
+    return func
+
+
+def check_type_not_field(form_name, key, model, field):
+    def func(self):
+        return not check_type_field(form_name, key, model, field)(self)
+    return func
+
+
 def check_treatment(form_name, type_key, type_list=[], not_type_list=[]):
     type_list = [models.TreatmentType.objects.get(txt_idx=tpe).pk
                  for tpe in type_list]
@@ -299,8 +351,8 @@ def check_treatment(form_name, type_key, type_list=[], not_type_list=[]):
            request.session[storage.prefix]['step_data'][form_name]:
             return False
         try:
-            tpe = request.session[storage.prefix]['step_data']\
-                [form_name][form_name + '-' + type_key]
+            tpe = request.session[storage.prefix][
+                'step_data'][form_name][form_name + '-' + type_key]
             if not tpe:
                 return False
             type = int(tpe[0])
@@ -342,6 +394,10 @@ class FindDeletionForm(FinalForm):
 class UpstreamFindFormSelection(FindFormSelection):
     form_label = _(u"Upstream find")
 
+    def __init__(self, *args, **kwargs):
+        super(UpstreamFindFormSelection, self).__init__(*args, **kwargs)
+        self.fields['pk'].required = True
+
 #############################################
 # Source management for archaelogical finds #
 #############################################
@@ -482,6 +538,21 @@ class BaseTreatmentForm(SelectFindBasketForm):
             exclude=['packaging'])
         self.fields['treatment_type'].help_text = \
             models.TreatmentType.get_help(exclude=['packaging'])
+        self.fields['basket'].required = False
+        self.fields['basket'].help_text = \
+            _(u"Leave it blank if you want to select a single item")
+
+    def clean(self, *args, **kwargs):
+        try:
+            treatment = models.TreatmentType.objects.get(
+                pk=self.cleaned_data['treatment_type'], available=True)
+        except models.TreatmentType.DoesNotExist:
+            raise forms.ValidationError(_(u"This treatment type is not "
+                                          u"available."))
+        if treatment.upstream_is_many and \
+                not self.cleaned_data.get('basket'):
+            raise forms.ValidationError(_(u"This treatment needs a basket."))
+        return self.cleaned_data
 
 """
 ####################################
diff --git a/archaeological_finds/models.py b/archaeological_finds/models.py
index 65d30f37a..f6483e989 100644
--- a/archaeological_finds/models.py
+++ b/archaeological_finds/models.py
@@ -613,10 +613,16 @@ class FindSource(Source):
 
 class TreatmentType(GeneralType):
     virtual = models.BooleanField(_(u"Virtual"))
-    upstream_is_many = models.BooleanField(_(u"Upstream is many"),
-                                           default=False)
-    downstream_is_many = models.BooleanField(_(u"Downstream is many"),
-                                             default=False)
+    upstream_is_many = models.BooleanField(
+        _(u"Upstream is many"), default=False,
+        help_text=_(
+            u"Check this if for this treatment from many finds you'll get "
+            u"one."))
+    downstream_is_many = models.BooleanField(
+        _(u"Downstream is many"), default=False,
+        help_text=_(
+            u"Check this if for this treatment from one find you'll get "
+            u"many."))
 
     class Meta:
         verbose_name = _(u"Treatment type")
@@ -633,8 +639,11 @@ class Treatment(BaseHistorizedItem, OwnPerms):
     comment = models.TextField(_(u"Comment"), blank=True, null=True)
     treatment_type = models.ForeignKey(TreatmentType,
                                        verbose_name=_(u"Treatment type"))
-    location = models.ForeignKey(Warehouse, verbose_name=_(u"Location"),
-                                 blank=True, null=True)
+    location = models.ForeignKey(
+        Warehouse, verbose_name=_(u"Location"), blank=True, null=True,
+        help_text=_(
+            u"Location where the treatment is done. Target warehouse for "
+            u"a move."))
     other_location = models.CharField(_(u"Other location"), max_length=200,
                                       blank=True, null=True)
     person = models.ForeignKey(
diff --git a/archaeological_finds/views.py b/archaeological_finds/views.py
index e55291a76..8bae9785c 100644
--- a/archaeological_finds/views.py
+++ b/archaeological_finds/views.py
@@ -319,26 +319,22 @@ class DeleteFindBasketView(IshtarMixin, LoginRequiredMixin, FormView):
 
 treatment_creation_wizard = TreatmentWizard.as_view([
     ('basetreatment-treatment_creation', BaseTreatmentForm),
-    # ('selecfind-treatment_creation', UpstreamFindFormSelection),
-    ('multiselecfinds-treatment_creation', FindMultipleFormSelection),
+    ('selecfind-treatment_creation', UpstreamFindFormSelection),
     ('resultfind-treatment_creation', ResultFindForm),
     ('resultfinds-treatment_creation', ResultFindFormSet),
     ('final-treatment_creation', FinalForm)],
     condition_dict={
     'selecfind-treatment_creation':
-        check_treatment('basetreatment-treatment_creation',
-                        'treatment_type', not_type_list=[
-                            'physical_grouping', 'packaging']),
-    'multiselecfinds-treatment_creation':
-        check_treatment('basetreatment-treatment_creation',
-                        'treatment_type', ['physical_grouping',
-                                           'packaging']),
+        check_not_exist('basetreatment-treatment_creation',
+                        'basket'),
     'resultfinds-treatment_creation':
-        check_treatment('basetreatment-treatment_creation',
-                        'treatment_type', ['split']),
+        check_type_field('basetreatment-treatment_creation',
+                         'treatment_type', models.TreatmentType,
+                         'downstream_is_many'),
     'resultfind-treatment_creation':
-        check_treatment('basetreatment-treatment_creation',
-                        'treatment_type', not_type_list=['split'])},
+        check_type_field('basetreatment-treatment_creation',
+                         'treatment_type', models.TreatmentType,
+                         'upstream_is_many')},
     label=_(u"New treatment"),
     url_name='treatment_creation',)