Manage main image: model, auto associate default image

1 files changed, 16 insertions, 2 deletions

diff --git a/archaeological_finds/models_treatments.py b/archaeological_finds/models_treatments.py
index 06f61068b..64bad203d 100644
--- a/archaeological_finds/models_treatments.py
+++ b/archaeological_finds/models_treatments.py
@@ -23,7 +23,8 @@ 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.db.models.signals import post_save, post_delete, pre_delete, \
+    m2m_changed
 from django.template.defaultfilters import slugify
 from django.utils.translation import ugettext_lazy as _, pgettext_lazy, \
     activate, deactivate
@@ -35,7 +36,7 @@ from archaeological_warehouse.models import Warehouse, Container
 from ishtar_common.models import Document, GeneralType, \
     ImageModel, BaseHistorizedItem, OwnPerms, HistoricalRecords, Person, \
     Organization, ValueGetter, post_save_cache, ShortMenuItem, \
-    DashboardFormItem, ExternalIdManager
+    DashboardFormItem, ExternalIdManager, document_attached_changed
 from ishtar_common.utils import cached_label_changed, get_current_year, \
     update_data
 
@@ -179,6 +180,9 @@ class Treatment(DashboardFormItem, ValueGetter, BaseHistorizedItem,
     documents = models.ManyToManyField(
         Document, related_name='treatments', verbose_name=_(u"Documents"),
         blank=True)
+    main_image = models.ForeignKey(
+        Document, related_name='main_image_treatments',
+        verbose_name=_(u"Main image"), blank=True, null=True)
     cached_label = models.TextField(_(u"Cached name"), null=True, blank=True,
                                     db_index=True)
     history = HistoricalRecords()
@@ -579,6 +583,9 @@ def pre_delete_treatment(sender, **kwargs):
 
 pre_delete.connect(pre_delete_treatment, sender=Treatment)
 
+m2m_changed.connect(document_attached_changed,
+                    sender=Treatment.documents.through)
+
 
 class AbsFindTreatments(models.Model):
     find = models.ForeignKey(Find, verbose_name=_(u"Find"),
@@ -847,6 +854,9 @@ class TreatmentFile(DashboardFormItem, ClosedItem, BaseHistorizedItem,
     documents = models.ManyToManyField(
         Document, related_name='treatment_files', verbose_name=_(u"Documents"),
         blank=True)
+    main_image = models.ForeignKey(
+        Document, related_name='main_image_treatment_files',
+        verbose_name=_(u"Main image"), blank=True, null=True)
     associated_basket = models.ForeignKey(
         FindBasket, null=True, blank=True, on_delete=models.SET_NULL,
         related_name='treatment_files'
@@ -956,4 +966,8 @@ class TreatmentFile(DashboardFormItem, ClosedItem, BaseHistorizedItem,
         super(TreatmentFile, self).save(*args, **kwargs)
 
 
+
+m2m_changed.connect(document_attached_changed,
+                    sender=TreatmentFile.documents.through)
+
 post_save.connect(cached_label_changed, sender=TreatmentFile)