Migrate images and sources to new document model (refs #4107)

1 files changed, 20 insertions, 14 deletions

diff --git a/archaeological_finds/models_treatments.py b/archaeological_finds/models_treatments.py
index 2baced6e1..19f00de1f 100644
--- a/archaeological_finds/models_treatments.py
+++ b/archaeological_finds/models_treatments.py
@@ -24,17 +24,16 @@ from django.contrib.gis.db import models
 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
-from django.utils.translation import ugettext_lazy as _, ugettext
+from django.utils.translation import ugettext_lazy as _
 
-
-from ishtar_common.utils import cached_label_changed, get_current_year
-from ishtar_common.models import GeneralType, ImageModel, BaseHistorizedItem, \
-    OwnPerms, HistoricalRecords, Person, Organization, Source, \
-    ValueGetter, post_save_cache, ShortMenuItem, DashboardFormItem, \
-    IshtarImage, ThroughImage
-from archaeological_warehouse.models import Warehouse, Container
 from archaeological_finds.models_finds import Find, FindBasket, TreatmentType
 from archaeological_operations.models import ClosedItem, Operation
+from archaeological_warehouse.models import Warehouse, Container
+from ishtar_common.models import BaseSource, Document, GeneralType, \
+    ImageModel, BaseHistorizedItem, OwnPerms, HistoricalRecords, Person, \
+    Organization, ValueGetter, post_save_cache, ShortMenuItem, \
+    DashboardFormItem, IshtarImage, ThroughImage
+from ishtar_common.utils import cached_label_changed, get_current_year
 
 
 class TreatmentState(GeneralType):
@@ -125,6 +124,9 @@ class Treatment(DashboardFormItem, ValueGetter, BaseHistorizedItem,
                                        blank=True, null=True)
     target_is_basket = models.BooleanField(_(u"Target a basket"),
                                            default=False)
+    documents = models.ManyToManyField(
+        Document, related_name='treatments', verbose_name=_(u"Documents"),
+        blank=True)
     images = models.ManyToManyField(IshtarImage, verbose_name=_(u"Images"),
                                     blank=True, through='TreatmentImage')
     cached_label = models.TextField(_(u"Cached name"), null=True, blank=True,
@@ -533,6 +535,9 @@ class TreatmentFile(DashboardFormItem, ClosedItem, BaseHistorizedItem,
     reception_date = models.DateField(_(u'Reception date'), blank=True,
                                       null=True)
     comment = models.TextField(_(u"Comment"), null=True, blank=True)
+    documents = models.ManyToManyField(
+        Document, related_name='treatment_files', verbose_name=_(u"Documents"),
+        blank=True)
     cached_label = models.TextField(_(u"Cached name"), null=True, blank=True,
                                     db_index=True)
     history = HistoricalRecords()
@@ -611,14 +616,14 @@ class TreatmentFile(DashboardFormItem, ClosedItem, BaseHistorizedItem,
 post_save.connect(cached_label_changed, sender=TreatmentFile)
 
 
-class TreatmentSource(Source):
+class TreatmentSource(BaseSource):
     treatment = models.ForeignKey(
         Treatment, verbose_name=_(u"Treatment"), related_name="source")
     BOOL_FIELDS = ['duplicate']
-    TABLE_COLS = ['treatment__cached_label'] + Source.TABLE_COLS
+    TABLE_COLS = ['treatment__cached_label'] + BaseSource.TABLE_COLS
     COL_LABELS = {'treatment__cached_label': _(u"Treatment")}
     SHOW_URL = 'show-treatmentsource'
-    PARENT_SEARCH_VECTORS = Source.PARENT_SEARCH_VECTORS + ['treatment']
+    PARENT_SEARCH_VECTORS = BaseSource.PARENT_SEARCH_VECTORS + ['treatment']
 
     class Meta:
         verbose_name = _(u"Treatment documentation")
@@ -644,15 +649,16 @@ class TreatmentSource(Source):
         return u"{}/source".format(self.treatment._get_base_image_path())
 
 
-class TreatmentFileSource(Source):
+class TreatmentFileSource(BaseSource):
     treatment_file = models.ForeignKey(
         TreatmentFile, verbose_name=_(u"Treatment request"),
         related_name="source")
     BOOL_FIELDS = ['duplicate']
-    TABLE_COLS = ['treatment_file__cached_label'] + Source.TABLE_COLS
+    TABLE_COLS = ['treatment_file__cached_label'] + BaseSource.TABLE_COLS
     COL_LABELS = {'treatment_file__cached_label': _(u"Treatment file")}
     SHOW_URL = 'show-treatmentfilesource'
-    PARENT_SEARCH_VECTORS = Source.PARENT_SEARCH_VECTORS + ['treatment_file']
+    PARENT_SEARCH_VECTORS = BaseSource.PARENT_SEARCH_VECTORS + [
+        'treatment_file']
 
     class Meta:
         verbose_name = _(u"Treatment request documentation")