Importer: let clients specify schemes of identifiers to preserve

mysociety · Jan 6, 2017 · ad332c1 · ad332c1
1 parent 6a2d14f
commit ad332c1
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 1 deletion.
diff --git a/popolo/importers/popolo_json.py b/popolo/importers/popolo_json.py
@@ -74,7 +74,7 @@ class PopoloJSONImporter(object):
 
     TRUNCATE_OPTIONS = set(['yes', 'warn', 'exception'])
 
-    def __init__(self, id_prefix=None, truncate='exception'):
+    def __init__(self, id_prefix=None, truncate='exception', **kwargs):
         super(PopoloJSONImporter, self).__init__()
         if id_prefix is None:
             self.id_prefix = 'popit-'
@@ -92,8 +92,19 @@ def __init__(self, id_prefix=None, truncate='exception'):
             'organization': {self.id_prefix + 'organization'},
             'area': {self.id_prefix + 'area'},
         }
+        # And if any extra ID schemes have been requested for
+        # preservation, keep them too:
+        if kwargs.get('id_schemes_to_preserve'):
+            self.add_id_schemes_to_preserve(kwargs['id_schemes_to_preserve'])
         self.observers = []
 
+    def add_id_schemes_to_preserve(self, id_schemes_to_preserve):
+        for collection, schemes in id_schemes_to_preserve.items():
+            if collection not in NEW_COLLECTIONS:
+                raise Exception("Unknown collection: '{}'".format(collection))
+            self.id_schemes_to_preserve.setdefault(collection, {})
+            self.id_schemes_to_preserve[collection].update(schemes)
+
     def add_observer(self, observer):
         self.observers.append(observer)
 

diff --git a/popolo/tests/test_importer.py b/popolo/tests/test_importer.py
@@ -1034,6 +1034,44 @@ def test_remove_identifiers_even_if_none_specified(self):
             [(u'popolo:person', u'a1b2')]
         )
 
+    def test_extra_preserved_identifiers(self):
+        existing_person = models.Person.objects.create(name='Algernon')
+        existing_person.identifiers.create(
+            scheme='popolo:person',
+            identifier="a1b2"
+        )
+        existing_person.identifiers.create(
+            scheme='ignorable',
+            identifier="some-data-we-care-nothing-for"
+        )
+        existing_person.identifiers.create(
+            scheme='preserve-me',
+            identifier="data-that-should-be-kept"
+        )
+        input_json = '''
+{
+    "persons": [
+        {
+            "id": "a1b2",
+            "name": "Alice",
+            "identifiers": []
+        }
+    ]
+}
+'''
+        data = json.loads(input_json)
+        importer = PopItImporter(
+            id_prefix='popolo:', id_schemes_to_preserve={'person': {'preserve-me'}})
+        importer.import_from_export_json_data(data)
+        self.assertEqual(models.Person.objects.count(), 1)
+        # Reget the person from the database:
+        person = models.Person.objects.get(pk=existing_person.id)
+        self.assertEqual(person.name, 'Alice')
+        self.assertSequenceEqual(
+            person.identifiers.order_by('scheme').values_list('scheme', 'identifier'),
+            [(u'popolo:person', u'a1b2'), (u'preserve-me', 'data-that-should-be-kept')]
+        )
+
 
 class SurprisingExceptionTests(TestCase):