Treatment file sheet: add treatment action

1 files changed, 20 insertions, 0 deletions

diff --git a/archaeological_finds/models_treatments.py b/archaeological_finds/models_treatments.py
index 1083b479b..15ef64040 100644
--- a/archaeological_finds/models_treatments.py
+++ b/archaeological_finds/models_treatments.py
@@ -21,6 +21,7 @@ import datetime
 
 from django.conf import settings
 from django.contrib.gis.db import models
+from django.core.urlresolvers import reverse
 from django.db.models import Max, Q
 from django.db.models.signals import post_save, post_delete, pre_delete
 from django.template.defaultfilters import slugify
@@ -699,6 +700,8 @@ class FindTreatments(AbsFindTreatments):
 
 
 class TreatmentFileType(GeneralType):
+    treatment_type = models.ForeignKey(TreatmentType, blank=True, null=True)
+
     class Meta:
         verbose_name = _(u"Treatment request type")
         verbose_name_plural = _(u"Treatment request types")
@@ -839,6 +842,23 @@ class TreatmentFile(DashboardFormItem, ClosedItem, BaseHistorizedItem,
                          for attr in ('year', 'index', 'internal_reference',
                                       'name') if getattr(self, attr)])
 
+    def get_extra_actions(self, request):
+        if not self.associated_basket:
+            return []
+        if self.type.treatment_type and self.treatments.filter(
+                treatment_types__pk=self.type.treatment_type.pk).count():
+            # a treatment of this type already exists
+            return []
+        actions = []
+        can_edit_find = self.can_do(request, 'change_find')
+        if can_edit_find:
+            actions += [
+                (reverse('treatmentfile-add-treatment', args=[self.pk]),
+                 _(u"Add associated treatment"), "fa fa-exchange", "", "",
+                 False),
+            ]
+        return actions
+
     @classmethod
     def get_owns(cls, user, menu_filtr=None, limit=None, values=None,
                  get_short_menu_class=None):