QA packaging: allow to change ref and current containers

1 files changed, 13 insertions, 1 deletions

diff --git a/archaeological_finds/models_treatments.py b/archaeological_finds/models_treatments.py
index 0a04cefb9..a8c48a3ad 100644
--- a/archaeological_finds/models_treatments.py
+++ b/archaeological_finds/models_treatments.py
@@ -434,7 +434,7 @@ class Treatment(DashboardFormItem, ValueGetter, BaseHistorizedItem,
     def save(self, *args, **kwargs):
         items, user, extra_args_for_new, resulting_find = [], None, [], None
         upstream_items, upstream_item, resulting_finds = [], None, None
-        treatment_types = []
+        treatment_types, return_new = [], False
         if "items" in kwargs:
             items = kwargs.pop('items')
         if "resulting_find" in kwargs:
@@ -451,6 +451,8 @@ class Treatment(DashboardFormItem, ValueGetter, BaseHistorizedItem,
             extra_args_for_new = kwargs.pop('extra_args_for_new')
         if "treatment_type_list" in kwargs:
             treatment_types = kwargs.pop('treatment_type_list')
+        if "return_new" in kwargs:
+            return_new = kwargs.pop('return_new')
         self.pre_save()
         super(Treatment, self).save(*args, **kwargs)
         to_be_executed = not self.executed and self.treatment_state.executed
@@ -491,6 +493,7 @@ class Treatment(DashboardFormItem, ValueGetter, BaseHistorizedItem,
         create_new_find = bool([tp for tp in treatment_types
                                 if tp.create_new_find])
 
+        new_items = []
         for item in items:
             if not create_new_find or not to_be_executed:
                 self.finds.add(item)
@@ -504,6 +507,7 @@ class Treatment(DashboardFormItem, ValueGetter, BaseHistorizedItem,
                     setattr(new, k, extra_args_for_new[k])
                 new.save()
                 updated.append(new.pk)
+                new_items.append(new)
                 # update baskets
                 for basket in \
                         FindBasket.objects.filter(items__pk=item.pk).all():
@@ -511,6 +515,8 @@ class Treatment(DashboardFormItem, ValueGetter, BaseHistorizedItem,
                     basket.items.add(new)
 
         if not to_be_executed:
+            if return_new:
+                return new_items
             return
 
         if create_new_find:
@@ -534,6 +540,8 @@ class Treatment(DashboardFormItem, ValueGetter, BaseHistorizedItem,
 
         # manage containers
         if not self.container:
+            if return_new:
+                return new_items
             return
 
         container_attrs = []
@@ -549,6 +557,8 @@ class Treatment(DashboardFormItem, ValueGetter, BaseHistorizedItem,
 
         if not container_attrs:
             # non consistent treatment
+            if return_new:
+                return new_items
             return
 
         for find in q.all():
@@ -561,6 +571,8 @@ class Treatment(DashboardFormItem, ValueGetter, BaseHistorizedItem,
                     find.save()
         self.executed = True
         self.save()
+        if return_new:
+            return new_items
 
     @property
     def associated_filename(self):