QA packaging: allow to change ref and current containers

1 files changed, 67 insertions, 20 deletions

diff --git a/archaeological_finds/forms_treatments.py b/archaeological_finds/forms_treatments.py
index 18c0697fc..5471cea23 100644
--- a/archaeological_finds/forms_treatments.py
+++ b/archaeological_finds/forms_treatments.py
@@ -357,15 +357,21 @@ class QAFindTreatmentForm(IshtarForm):
             reverse_lazy('autocomplete-container'),
             associated_model=Container, new=True),
         validators=[valid_id(Container)])
-    reference_container = forms.BooleanField(
-        label=_(u"Change the reference container"), required=False,
-        widget=widgets.CheckboxInput, initial=True,
-        help_text=_(u"If unchecked the current container will be changed")
+    container_to_change = forms.ChoiceField(
+        label=_(u"Change "), required=True,
+        choices=(
+            ('current-and-reference', _(u"current and reference containers")),
+            ('reference', _(u"the reference container")),
+            ('current', _(u"the current container")),
+        )
     )
     create_treatment = forms.BooleanField(
         label=_(u"Create a treatment"), required=False,
         widget=widgets.CheckboxInput
     )
+    treatment_type = forms.ChoiceField(
+        label=_(u"Treatment type"), choices=[], required=False
+    )
     year = forms.IntegerField(
         label=_("Year"), initial=lambda: datetime.datetime.now().year,
         validators=[validators.MinValueValidator(1000),
@@ -394,10 +400,36 @@ class QAFindTreatmentForm(IshtarForm):
                 self.user = self.user.ishtaruser
         self.items = kwargs.pop('items')
 
+        tt_change_ref_loca = list(
+            models.TreatmentType.objects.filter(
+                available=True, change_reference_location=True).all())
+
+        self.treatment_type_ref_choices = u"".join([
+            u"<option value='{}'>{}</option>".format(tt.pk, unicode(tt))
+            for tt in tt_change_ref_loca
+        ])
+
+        tt_change_current_loca = list(
+            models.TreatmentType.objects.filter(
+                available=True, change_current_location=True).all())
+
+        self.treatment_type_current_choices = u"".join([
+            u"<option value='{}'>{}</option>".format(tt.pk, unicode(tt))
+            for tt in tt_change_current_loca
+        ])
+        self.treatment_type_all_choices = self.treatment_type_ref_choices + \
+            self.treatment_type_current_choices
+
         super(QAFindTreatmentForm, self).__init__(*args, **kwargs)
         if not self.user:
             return
 
+        # treatment type is dynamic put all for check
+        self.fields['treatment_type'].choices = [
+            (tt.pk, unicode(tt)) for tt in (
+                    tt_change_ref_loca + tt_change_current_loca)
+        ]
+
         q = Person.objects.filter(ishtaruser__pk=self.user.pk)
         if q.count():
             person = q.all()[0]
@@ -413,16 +445,22 @@ class QAFindTreatmentForm(IshtarForm):
                 self.cleaned_data['year'] = self.cleaned_data['start_date'].year
             else:
                 raise forms.ValidationError(_(u"At least a year is required."))
+        if not self.cleaned_data.get('treatment_type', None):
+            raise forms.ValidationError(_(u"Treatment type is required."))
+
         return self.cleaned_data
 
     def save(self, items, user):
         container = Container.objects.get(pk=self.cleaned_data['container'])
+        container_to_change = self.cleaned_data.get('container_to_change', '')
         if self.cleaned_data['create_treatment']:
-            packaging, created = models.TreatmentType.objects.get_or_create(
-                txt_idx='packaging',
-                defaults={"label": _(u"Packaging"),
-                          "change_reference_location": True}
-            )
+            treat_type = models.TreatmentType.objects.get(
+                pk=self.cleaned_data['treatment_type'])
+            treat_state, __ = models.TreatmentState.objects.get_or_create(
+                txt_idx='completed',
+                defaults={
+                    'label': _(u"Completed"), 'executed': True,
+                    'available': True})
             t = models.Treatment.objects.create(
                 container=container,
                 year=self.cleaned_data['year'],
@@ -430,19 +468,28 @@ class QAFindTreatmentForm(IshtarForm):
                 location=container.location,
                 person_id=self.cleaned_data['person'],
                 organization_id=self.cleaned_data['organization'],
-                history_modifier=user
+                history_modifier=user,
+                treatment_state=treat_state
             )
-            t.treatment_types.add(packaging)
-            t.save(items=items)
-            return
-        container_attr = 'container'
-        if self.cleaned_data.get('reference_container', False):
-            container_attr = 'container_ref'
+            t.treatment_types.add(treat_type)
+            new_items = t.save(items=items, return_new=True)
+            if new_items:
+                items = new_items
+        container_attrs = []
+        if container_to_change in ('reference', 'current-and-reference'):
+            container_attrs.append('container_ref')
+        if container_to_change in ('current', 'current-and-reference'):
+            container_attrs.append('container')
         for find in items:
-            if getattr(find, container_attr) == container:
-                continue
-            setattr(find, container_attr, container)
-            find.save()
+            changed = False
+            for container_attr in container_attrs:
+                if getattr(find, container_attr) == container:
+                    continue
+                setattr(find, container_attr, container)
+                changed = True
+            if changed:
+                find.history_modifier = user
+                find.save()
 
 
 SLICING = (("month", _(u"months")), ('year', _(u"years")),)