Treatment n-1 - 1-n: fix treatment type filter

1 files changed, 7 insertions, 2 deletions

diff --git a/archaeological_finds/forms_treatments.py b/archaeological_finds/forms_treatments.py
index 53adc2813..10c4527ef 100644
--- a/archaeological_finds/forms_treatments.py
+++ b/archaeological_finds/forms_treatments.py
@@ -79,6 +79,8 @@ class TreatmentFormSelection(forms.Form):
 
 
 class BaseTreatmentForm(CustomForm, ManageOldType):
+    UPSTREAM_IS_MANY = False
+    DOWNSTREAM_IS_MANY = False
     form_label = _(u"Treatment")
     form_admin_name = _(u"Treatment - 020 - General")
     form_slug = "treatment-020-general"
@@ -205,8 +207,9 @@ class BaseTreatmentForm(CustomForm, ManageOldType):
         try:
             treatment_types = [
                 models.TreatmentType.objects.get(
-                    pk=pk, available=True, upstream_is_many=False,
-                    downstream_is_many=False)
+                    pk=pk, available=True,
+                    upstream_is_many=self.UPSTREAM_IS_MANY,
+                    downstream_is_many=self.DOWNSTREAM_IS_MANY)
                 for pk in data.get('treatment_type', [])]
         except models.TreatmentType.DoesNotExist:
             raise forms.ValidationError(_(u"Unknow treatment type"))
@@ -264,6 +267,7 @@ class BaseTreatmentForm(CustomForm, ManageOldType):
 
 
 class N1TreatmentForm(BaseTreatmentForm):
+    UPSTREAM_IS_MANY = True
     form_admin_name = _(u"Treatment n-1 - 020 - General")
     form_slug = "treatmentn1-020-general"
 
@@ -278,6 +282,7 @@ class N1TreatmentForm(BaseTreatmentForm):
 
 
 class OneNTreatmentForm(BaseTreatmentForm):
+    DOWNSTREAM_IS_MANY = True
     form_admin_name = _(u"Treatment 1-n - 020 - General")
     form_slug = "treatment1n-020-general"