Treatments: manage treatment without find duplication

1 files changed, 30 insertions, 16 deletions

diff --git a/archaeological_finds/models_treatments.py b/archaeological_finds/models_treatments.py
index 0bf3cad38..61d4ce5be 100644
--- a/archaeological_finds/models_treatments.py
+++ b/archaeological_finds/models_treatments.py
@@ -148,9 +148,9 @@ class Treatment(DashboardFormItem, ValueGetter, BaseHistorizedItem,
     start_date = models.DateField(_(u"Start date"), blank=True, null=True)
     end_date = models.DateField(_(u"Closing date"), blank=True, null=True)
     creation_date = models.DateTimeField(default=datetime.datetime.now)
-    find = models.ForeignKey(
-        "Find", verbose_name=_(u"Find"), related_name='treatments', blank=True,
-        null=True, help_text=_(u"Related find for non-destructive treatment"))
+    finds = models.ManyToManyField(
+        "Find", verbose_name=_(u"Finds"), related_name='treatments', blank=True,
+        help_text=_(u"Related finds for non-destructive treatment"))
     container = models.ForeignKey(Container, verbose_name=_(u"Container"),
                                   blank=True, null=True)
     estimated_cost = models.FloatField(_(u"Estimated cost"),
@@ -181,6 +181,7 @@ class Treatment(DashboardFormItem, ValueGetter, BaseHistorizedItem,
             ("change_own_treatment", u"Can change own Treatment"),
             ("delete_own_treatment", u"Can delete own Treatment"),
         )
+        ordering = ("start_date", )
 
     def __unicode__(self):
         if self.cached_label:
@@ -193,6 +194,10 @@ class Treatment(DashboardFormItem, ValueGetter, BaseHistorizedItem,
     def short_class_name(self):
         return _(u"TREATMENT")
 
+    @property
+    def limited_finds(self):
+        return self.finds.all()[:15]
+
     def natural_key(self):
         return (self.external_id, )
 
@@ -300,21 +305,30 @@ class Treatment(DashboardFormItem, ValueGetter, BaseHistorizedItem,
         updated = []
         if hasattr(items, "items"):
             items = items.items.all()
+
+        tps = list(self.treatment_types.all())
+        has_destructive_tp = bool([tp for tp in tps if tp.destructive])
         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()
-            updated.append(new.pk)
-            # update baskets
-            for basket in \
-                    FindBasket.objects.filter(items__pk=item.pk).all():
-                basket.items.remove(item)
-                basket.items.add(new)
+            if not has_destructive_tp:
+                self.finds.add(item)
+            else:
+                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()
+                updated.append(new.pk)
+                # update baskets
+                for basket in \
+                        FindBasket.objects.filter(items__pk=item.pk).all():
+                    basket.items.remove(item)
+                    basket.items.add(new)
+
         # manage containers
+        if not self.container:
+            return
         for find in Find.objects.filter(upstream_treatment=self).all():
             if find.container != self.container:
                 find.container = self.container