🚑️ document template: fix associated basket for treatment file and treatments

1 files changed, 19 insertions, 3 deletions

diff --git a/archaeological_finds/models_treatments.py b/archaeological_finds/models_treatments.py
index 60a4440fd..63ab10896 100644
--- a/archaeological_finds/models_treatments.py
+++ b/archaeological_finds/models_treatments.py
@@ -130,9 +130,24 @@ class AssociatedFindBasket:
             )
         )
 
+    def get_extra_values(self, prefix="", no_values=False, filtr=None, **kwargs):
+        values = {}
+        basket_key = prefix + "associated_basket_items"
+        if filtr and basket_key not in filtr:
+            return values
+        basket = self.associated_basket
+        if not basket:
+            return values
+        values[basket_key] = [
+            item.get_values(no_values=True, filtr=filtr, **kwargs)
+            for item in basket.items.distinct().all()
+        ]
+        return values
+
 
 class Treatment(
     DashboardFormItem,
+    AssociatedFindBasket,
     ValueGetter,
     DocumentItem,
     BaseHistorizedItem,
@@ -141,7 +156,6 @@ class Treatment(
     OwnPerms,
     ShortMenuItem,
     SheetItem,
-    AssociatedFindBasket,
 ):
     SLUG = "treatment"
     APP = "archaeological-finds"
@@ -483,7 +497,9 @@ class Treatment(
         return actions
 
     def get_extra_values(self, prefix="", no_values=False, filtr=None, **kwargs):
-        values = {}
+        values = super().get_extra_values(
+            prefix=prefix, no_values=False, filtr=None, **kwargs
+        )
         if not filtr or prefix + "upstream_finds" in filtr:
             values[prefix + "upstream_finds"] = " ; ".join(
                 str(up) for up in self.upstream.all()
@@ -1161,6 +1177,7 @@ post_delete.connect(post_save_cache, sender=TreatmentFileType)
 
 class TreatmentFile(
     DashboardFormItem,
+    AssociatedFindBasket,
     ClosedItem,
     DocumentItem,
     BaseHistorizedItem,
@@ -1168,7 +1185,6 @@ class TreatmentFile(
     OwnPerms,
     ValueGetter,
     MainItem,
-    AssociatedFindBasket,
 ):
     SLUG = "treatmentfile"
     APP = "archaeological-finds"