Manage fix of associated items (after imports or database dump)

3 files changed, 67 insertions, 2 deletions

diff --git a/archaeological_operations/models.py b/archaeological_operations/models.py
index 70c1c02ba..cf648a43a 100644
--- a/archaeological_operations/models.py
+++ b/archaeological_operations/models.py
@@ -38,7 +38,7 @@ from ishtar_common.models import GeneralType, BaseHistorizedItem, \
     SourceType, Person, Organization, Town, Dashboard, IshtarUser, ValueGetter,\
     DocumentTemplate, ShortMenuItem, DashboardFormItem, GeneralRelationType,\
     GeneralRecordRelations, post_delete_record_relation, OperationType, \
-    ImageModel, post_save_cache
+    ImageModel, post_save_cache, PersonType
 
 
 class RemainType(GeneralType):
@@ -252,6 +252,14 @@ class Operation(ClosedItem, BaseHistorizedItem, ImageModel, OwnPerms,
                            "common_name", "comment", "address", "old_code"]
     INT_SEARCH_VECTORS = ["year"]
     M2M_SEARCH_VECTORS = ["towns__name"]
+    ASSOCIATED = {
+        "scientist": {
+            ('person_types', PersonType): (
+                'head_scientist',
+                'sra_agent'
+            )
+        },
+    }
 
     # fields definition
     creation_date = models.DateField(_(u"Creation date"),
diff --git a/archaeological_operations/tests.py b/archaeological_operations/tests.py
index b75c02cae..ec7ae44c5 100644
--- a/archaeological_operations/tests.py
+++ b/archaeological_operations/tests.py
@@ -857,6 +857,23 @@ class OperationTest(TestCase, OperationInitTest):
         parcel.save()
         self.assertEqual(parcel.external_id, 'blabla')
 
+    def test_associated(self):
+        scientist = Person.objects.create(name="C-3PO")
+        self.item.scientist = scientist
+        self.item.save()
+        scientist = Person.objects.get(pk=scientist.pk)
+        self.assertIn(PersonType.objects.get(txt_idx='head_scientist'),
+                      scientist.person_types.all())
+
+        # do not change if in the list
+        sra = Person.objects.create(name="R2D2")
+        sra.person_types.add(PersonType.objects.get(txt_idx='sra_agent'))
+        self.item.scientist = sra
+        self.item.save()
+        self.assertNotIn(PersonType.objects.get(txt_idx='head_scientist'),
+                         sra.person_types.all())
+
+
     def create_relations(self):
         rel1 = models.RelationType.objects.create(
             symmetrical=True, label='Include', txt_idx='include')
diff --git a/ishtar_common/models.py b/ishtar_common/models.py
index 53c05d180..75f6adbed 100644
--- a/ishtar_common/models.py
+++ b/ishtar_common/models.py
@@ -1086,7 +1086,46 @@ class FullSearch(models.Model):
         return changed
 
 
-class BaseHistorizedItem(FullSearch, Imported, JsonData):
+class FixAssociated(object):
+    ASSOCIATED = {}
+
+    def fix_associated(self):
+        for key in self.ASSOCIATED:
+            item = getattr(self, key)
+            if not item:
+                continue
+            dct = self.ASSOCIATED[key]
+            for dct_key in dct:
+                subkey, ctype = dct_key
+                expected_values = dct[dct_key]
+                if not isinstance(expected_values, (list, tuple)):
+                    expected_values = [expected_values]
+                if hasattr(ctype, "txt_idx"):
+                    expected_values = [ctype.objects.get(txt_idx=v)
+                                       for v in expected_values]
+                current_vals = getattr(item, subkey)
+                is_many = False
+                if hasattr(current_vals, "all"):
+                    is_many = True
+                    current_vals = current_vals.all()
+                else:
+                    current_vals = [current_vals]
+                is_ok = False
+                for current_val in current_vals:
+                    if current_val in expected_values:
+                        is_ok = True
+                        break
+                if is_ok:
+                    continue
+                # the first value is used
+                new_value = expected_values[0]
+                if is_many:
+                    getattr(item, subkey).add(new_value)
+                else:
+                    setattr(item, subkey, new_value)
+
+
+class BaseHistorizedItem(FullSearch, Imported, JsonData, FixAssociated):
     """
     Historized item with external ID management.
     All historized items are searcheable and have a data json field
@@ -1267,6 +1306,7 @@ class BaseHistorizedItem(FullSearch, Imported, JsonData):
         for dep in self.EXTERNAL_ID_DEPENDENCIES:
             for obj in getattr(self, dep).all():
                 obj.update_external_id(save=True)
+        self.fix_associated()
         return True