Treatment: manage treatment with baskets

4 files changed, 76 insertions, 50 deletions

diff --git a/archaeological_finds/forms.py b/archaeological_finds/forms.py
index 271ee0c9d..67eb441e5 100644
--- a/archaeological_finds/forms.py
+++ b/archaeological_finds/forms.py
@@ -350,15 +350,23 @@ class FindMultipleFormSelection(forms.Form):
         return self.cleaned_data
 
 
+def check_form(wizard, form_name, key):
+    request = wizard.request
+    storage = wizard.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
+    return True
+
+
 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]:
+        if not check_form(self, form_name, key):
             return False
         try:
             val = request.session[storage.prefix][
@@ -377,15 +385,28 @@ def check_not_exist(form_name, key):
     return func
 
 
+def check_value(form_name, key, value):
+    def func(self):
+        request = self.request
+        storage = self.storage
+        if not check_form(self, form_name, key):
+            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 unicode(val) == unicode(value)
+        except ValueError:
+            return False
+    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]:
+        if not check_form(self, form_name, key):
             return False
         try:
             val = model.objects.get(pk=request.session[storage.prefix][
@@ -411,11 +432,7 @@ def check_treatment(form_name, type_key, type_list=[], not_type_list=[]):
     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 + '-' + type_key not in \
-           request.session[storage.prefix]['step_data'][form_name]:
+        if not check_form(self, form_name, type_key):
             return False
         try:
             tpe = request.session[storage.prefix][
@@ -534,6 +551,10 @@ class NewFindBasketForm(forms.ModelForm):
 
 
 class SelectFindBasketForm(forms.Form):
+    form_label = _(u"Basket")
+    associated_models = {'basket': models.FindBasket}
+    need_user_for_initialization = True
+
     basket = forms.ChoiceField(label=_(u"Basket"), required=True, choices=[])
 
     def __init__(self, *args, **kwargs):
@@ -595,7 +616,7 @@ class TreatmentFormSelection(forms.Form):
         validators=[valid_id(models.Treatment)])
 
 
-class BaseTreatmentForm(ManageOldType, SelectFindBasketForm):
+class BaseTreatmentForm(ManageOldType, forms.Form):
     form_label = _(u"Base treatment")
     base_models = ['treatment_type']
     associated_models = {'treatment_type': models.TreatmentType,
@@ -603,7 +624,7 @@ class BaseTreatmentForm(ManageOldType, SelectFindBasketForm):
                          'location': Warehouse,
                          'organization': Organization,
                          'container': models.Container,
-                         'basket': models.FindBasket}
+                         }
     need_user_for_initialization = True
 
     label = forms.CharField(label=_(u"Label"),
@@ -682,7 +703,6 @@ class BaseTreatmentForm(ManageOldType, SelectFindBasketForm):
             models.TreatmentType.get_help(
                dct={'upstream_is_many': False, 'downstream_is_many': False})
         # TODO
-        self.fields.pop('basket')
         """
         self.fields['basket'].required = False
         self.fields['basket'].help_text = \
diff --git a/archaeological_finds/models.py b/archaeological_finds/models.py
index 01c904cc8..3e0d09ef8 100644
--- a/archaeological_finds/models.py
+++ b/archaeological_finds/models.py
@@ -879,27 +879,25 @@ class Treatment(BaseHistorizedItem, ImageModel, OwnPerms):
             user = kwargs.pop('user')
         if "extra_args_for_new" in kwargs:
             extra_args_for_new = kwargs.pop('extra_args_for_new')
-        is_new = self.pk is None
         self.pre_save()
         super(Treatment, self).save(*args, **kwargs)
         updated = []
-        if is_new:
-            if hasattr(items, "items"):
-                items = items.items.all()
-            for item in items:
-                new = item.duplicate(user)
-                item.downstream_treatment = self
-                item.save()
-                new.upstream_treatment = self
-                for k in extra_args_for_new:
-                    setattr(new, k, extra_args_for_new[k])
-                new.save()
-                updated.append(new.pk)
-                # update baskets
-                for basket in \
-                        FindBasket.objects.filter(items__pk=item.pk).all():
-                    basket.items.remove(item)
-                    basket.items.add(new)
+        if hasattr(items, "items"):
+            items = items.items.all()
+        for item in items:
+            new = item.duplicate(user)
+            item.downstream_treatment = self
+            item.save()
+            new.upstream_treatment = self
+            for k in extra_args_for_new:
+                setattr(new, k, extra_args_for_new[k])
+            new.save()
+            updated.append(new.pk)
+            # update baskets
+            for basket in \
+                    FindBasket.objects.filter(items__pk=item.pk).all():
+                basket.items.remove(item)
+                basket.items.add(new)
         # manage containers
         for find in Find.objects.filter(upstream_treatment=self).all():
             if find.container != self.container:
diff --git a/archaeological_finds/views.py b/archaeological_finds/views.py
index 1f3213504..29d4f7da9 100644
--- a/archaeological_finds/views.py
+++ b/archaeological_finds/views.py
@@ -292,8 +292,9 @@ get_downstreamtreatment = get_item(
 treatment_wizard_steps = [
     ('basetreatment-treatment_creation', BaseTreatmentForm),
     ('selecfind-treatment_creation', UpstreamFindFormSelection),
-    ('resultfind-treatment_creation', ResultFindForm),
-    ('resultfinds-treatment_creation', ResultFindFormSet),
+    ('selecbasket-treatment_creation', SelectFindBasketForm),
+    # ('resultfind-treatment_creation', ResultFindForm),
+    # ('resultfinds-treatment_creation', ResultFindFormSet),
     ('final-treatment_creation', FinalForm)]
 
 treatment_search_wizard = SearchWizard.as_view([
@@ -305,16 +306,20 @@ treatment_creation_wizard = TreatmentWizard.as_view(
     treatment_wizard_steps,
     condition_dict={
     'selecfind-treatment_creation':
-        check_not_exist('basetreatment-treatment_creation',
-                        'basket'),
-    'resultfinds-treatment_creation':
-        check_type_field('basetreatment-treatment_creation',
-                         'treatment_type', models.TreatmentType,
-                         'downstream_is_many'),
-    'resultfind-treatment_creation':
-        check_type_field('basetreatment-treatment_creation',
-                         'treatment_type', models.TreatmentType,
-                         'upstream_is_many')},
+        check_value('basetreatment-treatment_creation',
+                    'target_is_basket', False),
+    'selecbasket-treatment_creation':
+        check_value('basetreatment-treatment_creation',
+                    'target_is_basket', True),
+    # 'resultfinds-treatment_creation':
+    #     check_type_field('basetreatment-treatment_creation',
+    #                      'treatment_type', models.TreatmentType,
+    #                      'downstream_is_many'),
+    # 'resultfind-treatment_creation':
+    #     check_type_field('basetreatment-treatment_creation',
+    #                      'treatment_type', models.TreatmentType,
+    #                      'upstream_is_many')
+    },
     label=_(u"New treatment"),
     url_name='treatment_creation',)
 
diff --git a/archaeological_finds/wizards.py b/archaeological_finds/wizards.py
index 9c91ac516..962f21e69 100644
--- a/archaeological_finds/wizards.py
+++ b/archaeological_finds/wizards.py
@@ -107,10 +107,13 @@ class TreatmentWizard(Wizard):
                         and not find.is_own(dct['history_modifier']):
                     raise PermissionDenied
                 dct['items'] = [find]
-            except (models.Find.DoesNotExist):
+            except models.Find.DoesNotExist:
                 raise PermissionDenied
         if 'basket' in dct:
-            dct.pop('basket')
+            basket = dct.pop('basket')
+            if basket.user.pk != dct['history_modifier'].pk:
+                raise PermissionDenied
+            dct['items'] = list(basket.items.all())
         return dct