Manage 1<->n treatments wizard

1 files changed, 44 insertions, 4 deletions

diff --git a/archaeological_finds/models_treatments.py b/archaeological_finds/models_treatments.py
index d92c98235..68935f142 100644
--- a/archaeological_finds/models_treatments.py
+++ b/archaeological_finds/models_treatments.py
@@ -29,6 +29,7 @@ from django.utils.translation import ugettext_lazy as _, pgettext_lazy, \
 
 from archaeological_finds.models_finds import Find, FindBasket, TreatmentType
 from archaeological_operations.models import ClosedItem, Operation
+from archaeological_context_records.models import Dating
 from archaeological_warehouse.models import Warehouse, Container
 from ishtar_common.models import Document, GeneralType, \
     ImageModel, BaseHistorizedItem, OwnPerms, HistoricalRecords, Person, \
@@ -72,7 +73,8 @@ class Treatment(DashboardFormItem, ValueGetter, BaseHistorizedItem,
         "person__cached_label": _(u"Responsible"),
     }
     # extra keys than can be passed to save method
-    EXTRA_SAVED_KEYS = ('items', 'user', 'resulting_find', 'upstream_items')
+    EXTRA_SAVED_KEYS = ('items', 'user', 'resulting_find', 'upstream_items',
+                        'resulting_finds', 'upstream_item')
 
     # alternative names of fields for searches
     ALT_NAMES = {
@@ -293,7 +295,10 @@ 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):
+    def _create_n_1_resulting_find(self, resulting_find, upstream_items):
+        """
+        Manage creation of n<->1 treatment
+        """
         m2m = {}
         base_fields = [f.name for f in Find._meta.get_fields()]
         for k in resulting_find.keys():
@@ -363,15 +368,45 @@ class Treatment(DashboardFormItem, ValueGetter, BaseHistorizedItem,
         new_find.skip_history_when_saving = True
         new_find.save()
 
+    def _create_1_n_resulting_find(self, resulting_finds, upstream_item, user):
+        """
+        Manage creation of 1<->n treatment
+        """
+        new_items = []
+        start_number = resulting_finds['start_number']
+        for idx in range(resulting_finds['number']):
+            label = resulting_finds['label'] + unicode(start_number + idx)
+            new_find = Find.objects.get(
+                pk=upstream_item.pk).duplicate(user)
+            new_find.upstream_treatment = self
+            new_find.label = label
+            new_find.skip_history_when_saving = True
+            new_find.save()
+            new_items.append(new_find)
+
+        upstream_item.downstream_treatment = self
+        upstream_item.skip_history_when_saving = True
+        upstream_item.save()
+
+        if getattr(user, 'ishtaruser', None):
+            b = FindBasket.objects.create(
+                label=resulting_finds['basket_name'], user=user.ishtaruser)
+            for item in new_items:
+                b.items.add(item)
+
     def save(self, *args, **kwargs):
         items, user, extra_args_for_new, resulting_find = [], None, [], None
-        upstream_items = []
+        upstream_items, upstream_item, resulting_finds = [], None, None
         if "items" in kwargs:
             items = kwargs.pop('items')
         if "resulting_find" in kwargs:
             resulting_find = kwargs.pop('resulting_find')
+        if "resulting_finds" in kwargs:
+            resulting_finds = kwargs.pop('resulting_finds')
         if "upstream_items" in kwargs:
             upstream_items = kwargs.pop('upstream_items')
+        if "upstream_item" in kwargs:
+            upstream_item = kwargs.pop('upstream_item')
         if "user" in kwargs:
             user = kwargs.pop('user')
         if "extra_args_for_new" in kwargs:
@@ -386,7 +421,12 @@ class Treatment(DashboardFormItem, ValueGetter, BaseHistorizedItem,
             upstream_items = upstream_items.items.all()
 
         if upstream_items and resulting_find:
-            self._create_resulting_find(resulting_find, upstream_items)
+            self._create_n_1_resulting_find(resulting_find, upstream_items)
+            return
+
+        if upstream_item and resulting_finds:
+            self._create_1_n_resulting_find(resulting_finds, upstream_item,
+                                            self.history_modifier)
             return
 
         tps = list(self.treatment_types.all())