QA packaging: can associate collection

1 files changed, 14 insertions, 4 deletions

diff --git a/archaeological_finds/forms_treatments.py b/archaeological_finds/forms_treatments.py
index b80b6c0fe..c62dbcc57 100644
--- a/archaeological_finds/forms_treatments.py
+++ b/archaeological_finds/forms_treatments.py
@@ -399,25 +399,29 @@ class TreatmentDeletionForm(FinalForm):
 
 class QAFindTreatmentForm(IshtarForm):
     container = forms.IntegerField(
-        label=_(u"Container"),
+        label=_("Container"),
         widget=widgets.JQueryAutoComplete(
             reverse_lazy('autocomplete-container'),
             associated_model=Container, new=True),
         validators=[valid_id(Container)])
     container_to_change = forms.ChoiceField(
-        label=_(u"Change "), required=True,
+        label=_("Change "), required=True,
         choices=(
             ('current-and-reference', _(u"current and reference containers")),
             ('reference', _(u"the reference container")),
             ('current', _(u"the current container")),
         )
     )
+    collection = forms.BooleanField(
+        label=_("Associate to the collection of the container"), required=False,
+        widget=widgets.CheckboxInput
+    )
     create_treatment = forms.BooleanField(
-        label=_(u"Create a treatment"), required=False,
+        label=_("Create a treatment"), required=False,
         widget=widgets.CheckboxInput
     )
     treatment_type = forms.ChoiceField(
-        label=_(u"Treatment type"), choices=[], required=False
+        label=_("Treatment type"), choices=[], required=False
     )
     year = forms.IntegerField(
         label=_("Year"), initial=lambda: datetime.datetime.now().year,
@@ -527,8 +531,14 @@ class QAFindTreatmentForm(IshtarForm):
             container_attrs.append('container_ref')
         if container_to_change in ('current', 'current-and-reference'):
             container_attrs.append('container')
+        collection = None
+        if self.cleaned_data.get("collection"):
+            collection = container.location_id
         for find in items:
             changed = False
+            if collection and find.collection_id != collection:
+                find.collection_id = collection
+                changed = True
             for container_attr in container_attrs:
                 if getattr(find, container_attr) == container:
                     continue