Add remains and periods to archaeological sites (refs #2251)

1 files changed, 19 insertions, 2 deletions

diff --git a/archaeological_operations/forms.py b/archaeological_operations/forms.py
index f411b20d0..0a8f9d94d 100644
--- a/archaeological_operations/forms.py
+++ b/archaeological_operations/forms.py
@@ -1,6 +1,6 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
-# Copyright (C) 2010-2014  Étienne Loks  <etienne.loks_AT_peacefrogsDOTnet>
+# Copyright (C) 2010-2015  É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
@@ -644,12 +644,22 @@ PeriodFormset.form_label = _("Periods")
 class ArchaeologicalSiteForm(forms.Form):
     reference = forms.CharField(label=_(u"Reference"), max_length=20)
     name = forms.CharField(label=_(u"Name"), max_length=200, required=False)
+    periods = forms.MultipleChoiceField(label=_("Periods"),
+                             choices=[], widget=forms.CheckboxSelectMultiple)
+    remains = forms.MultipleChoiceField(label=_("Remains"),
+                             choices=[], widget=forms.CheckboxSelectMultiple)
 
     def __init__(self, *args, **kwargs):
         self.limits = {}
         if 'limits' in kwargs:
             limits = kwargs.pop('limits')
         super(ArchaeologicalSiteForm, self).__init__(*args, **kwargs)
+        self.fields['periods'].choices = \
+                            models.Period.get_types(empty_first=False)
+        self.fields['periods'].help_text = models.Period.get_help()
+        self.fields['remains'].choices = \
+                            models.RemainType.get_types(empty_first=False)
+        self.fields['remains'].help_text = models.RemainType.get_help()
 
     def clean_reference(self):
         reference = self.cleaned_data['reference']
@@ -661,7 +671,14 @@ class ArchaeologicalSiteForm(forms.Form):
     def save(self, user):
         dct = self.cleaned_data
         dct['history_modifier'] = user
-        return models.ArchaeologicalSite.objects.create(**dct)
+        periods = dct.pop('periods')
+        remains = dct.pop('remains')
+        item = models.ArchaeologicalSite.objects.create(**dct)
+        for period in periods:
+            item.periods.add(period)
+        for remain in remains:
+            item.remains.add(remain)
+        return item
 
 class ArchaeologicalSiteSelectionForm(forms.Form):
     form_label = _("Associated archaelogical sites")