Manage reference containers - manage loan and loan return

1 files changed, 35 insertions, 5 deletions

diff --git a/archaeological_finds/models_treatments.py b/archaeological_finds/models_treatments.py
index cc0d17dd0..1083b479b 100644
--- a/archaeological_finds/models_treatments.py
+++ b/archaeological_finds/models_treatments.py
@@ -458,16 +458,46 @@ class Treatment(DashboardFormItem, ValueGetter, BaseHistorizedItem,
                     basket.items.remove(item)
                     basket.items.add(new)
 
-        # manage containers
-        if not self.container:
-            return
         if create_new_find:
             q = Find.objects.filter(upstream_treatment=self)
         else:
             q = Find.objects.filter(treatments=self)
+
+        # manage loan return
+        for tp in treatment_types:
+            if tp.txt_idx == 'loan-return':
+                for find in q.all():
+                    if find.container_ref:
+                        find.container = find.container_ref
+                        if find.pk in updated:
+                            # don't record twice history
+                            find.skip_history_when_saving = True
+                        find.save()
+
+        # manage containers
+        if not self.container:
+            return
+
+        container_attr = None
+        for tp in treatment_types:
+            if tp.txt_idx == 'loan':
+                if container_attr:
+                    # non consistent treatment
+                    return
+                container_attr = 'container'
+            if tp.txt_idx == 'packaging':
+                if container_attr:
+                    # non consistent treatment
+                    return
+                container_attr = 'container_ref'
+
+        if not container_attr:
+            # non consistent treatment
+            return
+
         for find in q.all():
-            if find.container != self.container:
-                find.container = self.container
+            if getattr(find, container_attr) != self.container:
+                setattr(find, container_attr, self.container)
                 if find.pk in updated:
                     # don't record twice history
                     find.skip_history_when_saving = True