1 files changed, 30 insertions, 2 deletions
diff --git a/archaeological_finds/models_finds.py b/archaeological_finds/models_finds.py
index 68d49dfd6..fa0a9df9d 100644
--- a/archaeological_finds/models_finds.py
+++ b/archaeological_finds/models_finds.py
@@ -71,7 +71,7 @@ from ishtar_common.models import (
     SearchVectorConfig,
     ValueGetter,
 )
-from ishtar_common.models_common import HistoricalRecords, SerializeItem, \
+from ishtar_common.models_common import HistoricalRecords, Imported, SerializeItem, \
     GeoVectorData, geodata_attached_changed
 from ishtar_common.utils import PRIVATE_FIELDS
 
@@ -2476,6 +2476,31 @@ class Find(
     def has_packaging_for_current_container(self):
         return FindTreatment.objects.filter(find=self, location_type__in=["B", "C"])
 
+    @staticmethod
+    def _get_upstream_count(upstream_id, idx):
+        if not upstream_id:
+            return idx
+        q = Find.objects.filter(downstream_treatment_id=upstream_id)
+        if not q.exists():  # upstream find has been deleted?
+            return idx
+        idx += 1
+        upstream_ids = q.values_list("upstream_treatment_id", flat=True)
+        return max([Find._get_upstream_count(up_id, idx) for up_id in upstream_ids])
+
+    @property
+    def upstream_count(self):
+        """
+        Count upstream treatments. Get the maximum count. If equal to 0 return None.
+        Used mainly to generate IDs
+        """
+        if not self.id:
+            return
+        if not hasattr(self, "_cache_upstream_count"):
+            self._cache_upstream_count = Find._get_upstream_count(
+                self.upstream_treatment_id, 0
+            ) or None
+        return self._cache_upstream_count
+
     @property
     def has_packaging_for_reference_container(self):
         return FindTreatment.objects.filter(find=self, location_type__in=["B", "R"])
@@ -3717,7 +3742,7 @@ LOCATION_TYPE = [
 LOCATION_TYPE_DICT = dict(LOCATION_TYPE)
 
 
-class FindTreatment(models.Model):
+class FindTreatment(Imported):
     """
     Record all new location for a find.
     """
@@ -3740,6 +3765,9 @@ class FindTreatment(models.Model):
         verbose_name_plural = _("Find - Treatments")
         db_table = 'archaeological_finds_find_treatments'
 
+    def __str__(self):
+        return f"{self.treatment} / {self.find}"
+
     def location_type_label(self):
         if self.location_type in LOCATION_TYPE_DICT:
             return LOCATION_TYPE_DICT[self.location_type]