Treatment n<->1: manage saving

1 files changed, 72 insertions, 2 deletions

diff --git a/archaeological_finds/models_treatments.py b/archaeological_finds/models_treatments.py
index 160e9d32d..4e2eb0faf 100644
--- a/archaeological_finds/models_treatments.py
+++ b/archaeological_finds/models_treatments.py
@@ -71,7 +71,7 @@ class Treatment(DashboardFormItem, ValueGetter, BaseHistorizedItem,
         "person__cached_label": _(u"Responsible"),
     }
     # extra keys than can be passed to save method
-    EXTRA_SAVED_KEYS = ('items', 'user')
+    EXTRA_SAVED_KEYS = ('items', 'user', 'resulting_find', 'upstream_items')
 
     # alternative names of fields for searches
     ALT_NAMES = {
@@ -292,10 +292,73 @@ class Treatment(DashboardFormItem, ValueGetter, BaseHistorizedItem,
         if q.count():
             self.index = q.all().aggregate(Max('index'))['index__max'] + 1
 
+    def _create_resulting_find(self, resulting_find, upstream_items):
+        m2m = {}
+        base_fields = [f.name for f in Find._meta.get_fields()]
+        for k in resulting_find.keys():
+            # if not in base fields should be a m2m
+            if k not in base_fields:
+                values = resulting_find.pop(k)
+                if values:
+                    m2m[k + "s"] = values
+
+        resulting_find['history_modifier'] = self.history_modifier
+        new_find = Find.objects.create(**resulting_find)
+
+        for k in m2m:
+            m2m_field = getattr(new_find, k)
+            try:
+                for value in m2m[k]:
+                    m2m_field.add(value)
+            except TypeError:
+                m2m_field.add(m2m[k])
+
+        dating_keys = ["period", "start_date", "end_date", "dating_type",
+                       "quality", "precise_dating"]
+        current_datings = []
+        current_base_finds = []
+        for upstream_item in upstream_items:
+            # datings are not explicitly part of the resulting_find
+            # need to reassociate with no duplicate
+            for dating in upstream_item.datings.all():
+                current_dating = []
+                for key in dating_keys:
+                    value = getattr(dating, key)
+                    if hasattr(value, 'pk'):
+                        value = value.pk
+                    current_dating.append(value)
+                current_dating = tuple(current_dating)
+                if current_dating in current_datings:
+                    # do not add similar dating
+                    continue
+                current_datings.append(current_dating)
+                dating.pk = None  # duplicate
+                dating.save()
+                new_find.datings.add(dating)
+
+            # associate base finds
+            for base_find in upstream_item.base_finds.all():
+                if base_find.pk in current_base_finds:
+                    continue
+                current_base_finds.append(base_find.pk)
+                new_find.base_finds.add(base_find)
+
+            upstream_item.downstream_treatment = self
+            upstream_item.history_modifier = self.history_modifier
+            upstream_item.save()
+        new_find.upstream_treatment = self
+        new_find.skip_history_when_saving = True
+        new_find.save()
+
     def save(self, *args, **kwargs):
-        items, user, extra_args_for_new = [], None, []
+        items, user, extra_args_for_new, resulting_find = [], None, [], None
+        upstream_items = []
         if "items" in kwargs:
             items = kwargs.pop('items')
+        if "resulting_find" in kwargs:
+            resulting_find = kwargs.pop('resulting_find')
+        if "upstream_items" in kwargs:
+            upstream_items = kwargs.pop('upstream_items')
         if "user" in kwargs:
             user = kwargs.pop('user')
         if "extra_args_for_new" in kwargs:
@@ -303,8 +366,15 @@ class Treatment(DashboardFormItem, ValueGetter, BaseHistorizedItem,
         self.pre_save()
         super(Treatment, self).save(*args, **kwargs)
         updated = []
+        # baskets
         if hasattr(items, "items"):
             items = items.items.all()
+        if hasattr(upstream_items, "items"):
+            upstream_items = upstream_items.items.all()
+
+        if upstream_items and resulting_find:
+            self._create_resulting_find(resulting_find, upstream_items)
+            return
 
         tps = list(self.treatment_types.all())
         has_destructive_tp = bool([tp for tp in tps if tp.destructive])