Treatment: workon on basket update after treatment

1 files changed, 29 insertions, 0 deletions

diff --git a/archaeological_finds/models.py b/archaeological_finds/models.py
index c8b392310..0a1e59a68 100644
--- a/archaeological_finds/models.py
+++ b/archaeological_finds/models.py
@@ -612,6 +612,35 @@ class Treatment(BaseHistorizedItem, OwnPerms):
             lbl += u" %s %s" % (_(u"by"), unicode(self.person))
         return lbl
 
+    def save(self, *args, **kwargs):
+        items, user, extra_args_for_new = [], None, []
+        if "items" in kwargs:
+            items = kwargs.pop('items')
+        if "user" in kwargs:
+            user = kwargs.pop('user')
+        if "extra_args_for_new" in kwargs:
+            extra_args_for_new = kwargs.pop('extra_args_for_new')
+        is_new = self.pk is None
+        super(Treatment, self).save(*args, **kwargs)
+        if not is_new or not items:
+            return
+        basket = None
+        if hasattr(items, "items"):
+            basket = items
+            items = basket.items.all()
+        for item in items:
+            new = item.duplicate(user)
+            item.downstream_treatment = self
+            item.save()
+            new.upstream_treatment = self
+            for k in extra_args_for_new:
+                setattr(new, k, extra_args_for_new[k])
+            new.save()
+            # update baskets
+            for basket in FindBasket.objects.filter(items__pk=item.pk).all():
+                basket.items.remove(item)
+                basket.items.add(new)
+
 
 class TreatmentSource(Source):
     class Meta: