🗃️ database: add associated basket to treatments - is exhibition for treatments types

1 files changed, 36 insertions, 7 deletions

diff --git a/archaeological_finds/models_treatments.py b/archaeological_finds/models_treatments.py
index 935a14f2f..c7c4fcac9 100644
--- a/archaeological_finds/models_treatments.py
+++ b/archaeological_finds/models_treatments.py
@@ -96,6 +96,25 @@ post_save.connect(post_save_cache, sender=TreatmentState)
 post_delete.connect(post_save_cache, sender=TreatmentState)
 
 
+class AssociatedFindBasket:
+    @property
+    def associated_basket(self):
+        if not self.associated_basket_id:
+            return
+        try:
+            return FindBasket.objects.get(pk=self.associated_basket_id)
+        except FindBasket.DoesNotExist:
+            return
+
+    def pre_save_basket(self):
+        associated_basket = self.associated_basket
+        if self.associated_basket_id and not associated_basket:
+            self.associated_basket_id = None
+        elif associated_basket and associated_basket.pk != self.associated_basket_id:
+            self.associated_basket_id = associated_basket.pk
+
+
+
 class Treatment(
     DashboardFormItem,
     ValueGetter,
@@ -106,6 +125,7 @@ class Treatment(
     OwnPerms,
     ShortMenuItem,
     SheetItem,
+    AssociatedFindBasket,
 ):
     SLUG = "treatment"
     APP = "archaeological-finds"
@@ -308,6 +328,11 @@ class Treatment(
         blank=True,
         null=True,
     )
+    # prevent circular imports...
+    associated_basket_id = models.PositiveIntegerField(
+        verbose_name=_("Basket ID"), blank=True, null=True
+    )
+    # associated_basket = models.ForeignKey(FindBasket)
     history = HistoricalRecords(bases=[HistoryModel])
 
     class Meta:
@@ -629,6 +654,7 @@ class Treatment(
         return self._is_reference_container_changer
 
     def save(self, *args, **kwargs):
+        self.pre_save_basket()
         items, user, extra_args_for_new, resulting_find = [], None, [], None
         upstream_items, upstream_item, resulting_finds = [], None, None
         treatment_types, return_new = [], False
@@ -1104,6 +1130,7 @@ class TreatmentFileType(GeneralType):
     treatment_type = models.ForeignKey(
         TreatmentType, blank=True, null=True, on_delete=models.SET_NULL
     )
+    is_exhibition = models.BooleanField(_("Is an exhibition"), default=False)
 
     class Meta:
         verbose_name = _("Treatment request type")
@@ -1125,6 +1152,7 @@ class TreatmentFile(
     OwnPerms,
     ValueGetter,
     MainItem,
+    AssociatedFindBasket,
 ):
     SLUG = "treatmentfile"
     APP = "archaeological-finds"
@@ -1276,13 +1304,11 @@ class TreatmentFile(
         blank=True,
         null=True,
     )
-    associated_basket = models.ForeignKey(
-        FindBasket,
-        null=True,
-        blank=True,
-        on_delete=models.SET_NULL,
-        related_name="treatment_files",
+    # prevent circular imports...
+    associated_basket_id = models.PositiveIntegerField(
+        verbose_name=_("Basket ID"), blank=True, null=True
     )
+    # associated_basket = models.ForeignKey(FindBasket)
     history = HistoricalRecords()
 
     class Meta:
@@ -1303,6 +1329,9 @@ class TreatmentFile(
     def __str__(self):
         return self.cached_label or ""
 
+    def natural_key(self):
+        return (self.year, self.index)
+
     @property
     def short_class_name(self):
         return _("Treatment request")
@@ -1403,7 +1432,7 @@ class TreatmentFile(
 
     def save(self, *args, **kwargs):
         self.pre_save()
-        super(TreatmentFile, self).save(*args, **kwargs)
+        super().save(*args, **kwargs)
 
 
 m2m_changed.connect(document_attached_changed, sender=TreatmentFile.documents.through)