Reactivate treatments

8 files changed, 75 insertions, 73 deletions

diff --git a/archaeological_finds/forms.py b/archaeological_finds/forms.py
index f5387febe..9fdf2ddf2 100644
--- a/archaeological_finds/forms.py
+++ b/archaeological_finds/forms.py
@@ -265,39 +265,6 @@ class MultipleFindFormSelection(forms.Form):
         validators=[valid_id(models.Find)])
 
 
-class BaseTreatmentForm(forms.Form):
-    form_label = _(u"Base treatment")
-    associated_models = {'treatment_type': models.TreatmentType,
-                         'person': Person,
-                         'location': Warehouse}
-    treatment_type = forms.ChoiceField(label=_(u"Treatment type"), choices=[])
-    person = forms.IntegerField(
-        label=_(u"Person"),
-        widget=widgets.JQueryAutoComplete(
-            reverse_lazy('autocomplete-person'), associated_model=Person,
-            new=True),
-        validators=[valid_id(Person)])
-    location = forms.IntegerField(
-        label=_(u"Location"),
-        widget=widgets.JQueryAutoComplete(
-            reverse_lazy('autocomplete-warehouse'), associated_model=Warehouse,
-            new=True),
-        validators=[valid_id(Warehouse)])
-    description = forms.CharField(label=_(u"Description"),
-                                  widget=forms.Textarea, required=False)
-    start_date = forms.DateField(label=_(u"Start date"), required=False,
-                                 widget=widgets.JQueryDate)
-    end_date = forms.DateField(label=_(u"End date"), required=False,
-                               widget=widgets.JQueryDate)
-
-    def __init__(self, *args, **kwargs):
-        super(BaseTreatmentForm, self).__init__(*args, **kwargs)
-        self.fields['treatment_type'].choices = models.TreatmentType.get_types(
-            exclude=['packaging'])
-        self.fields['treatment_type'].help_text = \
-            models.TreatmentType.get_help(exclude=['packaging'])
-
-
 class FindMultipleFormSelection(forms.Form):
     form_label = _(u"Upstream finds")
     associated_models = {'finds': models.Find}
@@ -322,7 +289,9 @@ def check_treatment(form_name, type_key, type_list=[], not_type_list=[]):
     not_type_list = [models.TreatmentType.objects.get(txt_idx=tpe).pk
                      for tpe in not_type_list]
 
-    def func(self, request, storage):
+    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\
@@ -330,8 +299,11 @@ def check_treatment(form_name, type_key, type_list=[], not_type_list=[]):
            request.session[storage.prefix]['step_data'][form_name]:
             return False
         try:
-            type = int(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])
             return (not type_list or type in type_list) \
                 and type not in not_type_list
         except ValueError:
@@ -477,6 +449,40 @@ class FindBasketAddItemForm(forms.Form):
         basket.items.add(item)
         return basket
 
+
+class BaseTreatmentForm(SelectFindBasketForm):
+    form_label = _(u"Base treatment")
+    associated_models = {'treatment_type': models.TreatmentType,
+                         'person': Person,
+                         'location': Warehouse,
+                         'basket': models.FindBasket}
+    treatment_type = forms.ChoiceField(label=_(u"Treatment type"), choices=[])
+    person = forms.IntegerField(
+        label=_(u"Person"),
+        widget=widgets.JQueryAutoComplete(
+            reverse_lazy('autocomplete-person'), associated_model=Person,
+            new=True),
+        validators=[valid_id(Person)])
+    location = forms.IntegerField(
+        label=_(u"Location"),
+        widget=widgets.JQueryAutoComplete(
+            reverse_lazy('autocomplete-warehouse'), associated_model=Warehouse,
+            new=True),
+        validators=[valid_id(Warehouse)])
+    description = forms.CharField(label=_(u"Description"),
+                                  widget=forms.Textarea, required=False)
+    start_date = forms.DateField(label=_(u"Start date"), required=False,
+                                 widget=widgets.JQueryDate)
+    end_date = forms.DateField(label=_(u"End date"), required=False,
+                               widget=widgets.JQueryDate)
+
+    def __init__(self, *args, **kwargs):
+        super(BaseTreatmentForm, self).__init__(*args, **kwargs)
+        self.fields['treatment_type'].choices = models.TreatmentType.get_types(
+            exclude=['packaging'])
+        self.fields['treatment_type'].help_text = \
+            models.TreatmentType.get_help(exclude=['packaging'])
+
 """
 ####################################
 # Source management for treatments #
diff --git a/archaeological_finds/ishtar_menu.py b/archaeological_finds/ishtar_menu.py
index d7a67091a..f69f33f69 100644
--- a/archaeological_finds/ishtar_menu.py
+++ b/archaeological_finds/ishtar_menu.py
@@ -46,10 +46,11 @@ MENU_SECTIONS = [
                  model=models.Find,
                  access_controls=['change_find',
                                   'change_own_find']),
-             # MenuItem('treatment_creation', _(u"Add a treatment"),
-             #     model=models.Treatment,
-             #     access_controls=['add_treatment',
-             #                      'add_own_treatment']),
+             MenuItem(
+                 'treatment_creation', _(u"Add a treatment"),
+                 model=models.Treatment,
+                 access_controls=['change_find',
+                                  'change_own_find']),
              MenuItem(
                  'find_deletion', _(u"Deletion"),
                  model=models.Find,
diff --git a/archaeological_finds/urls.py b/archaeological_finds/urls.py
index 5c859a83f..4b28b0b2b 100644
--- a/archaeological_finds/urls.py
+++ b/archaeological_finds/urls.py
@@ -84,6 +84,8 @@ urlpatterns = patterns(
     url(r'^find_basket_deletion/$',
         check_rights(['change_find', 'change_own_find'])(
             views.DeleteFindBasketView.as_view()), name='delete_findbasket'),
+    url(r'treatment_creation/(?P<step>.+)?$',
+        views.treatment_creation_wizard, name='treatment_creation'),
 )
 
 urlpatterns += patterns(
diff --git a/archaeological_finds/views.py b/archaeological_finds/views.py
index d980ebd69..e55291a76 100644
--- a/archaeological_finds/views.py
+++ b/archaeological_finds/views.py
@@ -317,37 +317,30 @@ class DeleteFindBasketView(IshtarMixin, LoginRequiredMixin, FormView):
         form.save()
         return HttpResponseRedirect(self.get_success_url())
 
-"""
 treatment_creation_wizard = TreatmentWizard.as_view([
     ('basetreatment-treatment_creation', BaseTreatmentForm),
-    ('selecfind-treatment_creation', UpstreamFindFormSelection),
+    # ('selecfind-treatment_creation', UpstreamFindFormSelection),
     ('multiselecfinds-treatment_creation', FindMultipleFormSelection),
-    ('container-treatment_creation', ContainerForm),
     ('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']),
-        'resultfinds-treatment_creation':
-            check_treatment('basetreatment-treatment_creation',
-                            'treatment_type', ['split']),
-        'resultfind-treatment_creation':
-            check_treatment('basetreatment-treatment_creation',
-                            'treatment_type', not_type_list=['split']),
-        'container-treatment_creation':
-            check_treatment('basetreatment-treatment_creation',
-                            'treatment_type', ['packaging']),
-    },
+    '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']),
+    'resultfinds-treatment_creation':
+        check_treatment('basetreatment-treatment_creation',
+                        'treatment_type', ['split']),
+    'resultfind-treatment_creation':
+        check_treatment('basetreatment-treatment_creation',
+                        'treatment_type', not_type_list=['split'])},
     label=_(u"New treatment"),
     url_name='treatment_creation',)
-"""
 
 """
 treatment_source_creation_wizard = TreatmentSourceWizard.as_view([
diff --git a/archaeological_finds/wizards.py b/archaeological_finds/wizards.py
index fabe1039a..8a6b92440 100644
--- a/archaeological_finds/wizards.py
+++ b/archaeological_finds/wizards.py
@@ -85,6 +85,14 @@ class FindDeletionWizard(DeletionWizard):
 
 class TreatmentWizard(Wizard):
     model = models.Treatment
+    basket_step = 'basetreatment-treatment_creation'
+
+    def get_form_kwargs(self, step):
+        kwargs = super(TreatmentWizard, self).get_form_kwargs(step)
+        if self.basket_step not in step:
+            return kwargs
+        kwargs['user'] = self.request.user
+        return kwargs
 
 
 class FindSourceWizard(SourceWizard):
diff --git a/archaeological_warehouse/urls.py b/archaeological_warehouse/urls.py
index 58f594bcd..9094f0a91 100644
--- a/archaeological_warehouse/urls.py
+++ b/archaeological_warehouse/urls.py
@@ -1,6 +1,6 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
-# Copyright (C) 2010-2012 Étienne Loks  <etienne.loks_AT_peacefrogsDOTnet>
+# Copyright (C) 2010-2016 Étienne Loks  <etienne.loks_AT_peacefrogsDOTnet>
 
 # This program is free software: you can redistribute it and/or modify
 # it under the terms of the GNU Affero General Public License as
@@ -27,8 +27,6 @@ import views
 # forms
 urlpatterns = patterns(
     '',
-    # url(r'treatment_creation/(?P<step>.+)?$',
-    #    views.treatment_creation_wizard, name='treatment_creation'),
     url(r'warehouse_packaging/(?P<step>.+)?$',
         views.warehouse_packaging_wizard, name='warehouse_packaging'),
 )
diff --git a/archaeological_warehouse/views.py b/archaeological_warehouse/views.py
index be5e6150c..486c2a1e0 100644
--- a/archaeological_warehouse/views.py
+++ b/archaeological_warehouse/views.py
@@ -1,6 +1,6 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
-# Copyright (C) 2010-2012  Étienne Loks  <etienne.loks_AT_peacefrogsDOTnet>
+# Copyright (C) 2010-2016  Étienne Loks  <etienne.loks_AT_peacefrogsDOTnet>
 
 # This program is free software: you can redistribute it and/or modify
 # it under the terms of the GNU Affero General Public License as
diff --git a/archaeological_warehouse/wizards.py b/archaeological_warehouse/wizards.py
index 409f7b28f..62b441e00 100644
--- a/archaeological_warehouse/wizards.py
+++ b/archaeological_warehouse/wizards.py
@@ -25,13 +25,7 @@ from archaeological_finds.models import Treatment
 
 
 class PackagingWizard(TreatmentWizard):
-
-    def get_form_kwargs(self, step):
-        kwargs = super(PackagingWizard, self).get_form_kwargs(step)
-        if 'base-packaging' not in step:
-            return kwargs
-        kwargs['user'] = self.request.user
-        return kwargs
+    basket_step = 'base-packaging'
 
     def save_model(self, dct, m2m, whole_associated_models, form_list,
                    return_object):