Fix division importer

1 files changed, 18 insertions, 10 deletions

diff --git a/archaeological_warehouse/models.py b/archaeological_warehouse/models.py
index 71751fcbe..3b319881d 100644
--- a/archaeological_warehouse/models.py
+++ b/archaeological_warehouse/models.py
@@ -386,11 +386,11 @@ class Warehouse(Address, DocumentItem, GeoItem, CompleteIdentifierItem,
         return current_container_type, previous_container_types
 
     @post_importer_action
-    def add_localisations(self, __, value):
-        self._add_localisations(value)
+    def add_localisations(self, context, value):
+        self._add_localisations(context, value)
     add_localisations.post_save = True
 
-    def _add_localisations(self, value, return_errors=False):
+    def _add_localisations(self, context, value, return_errors=False):
         """
         Add localisations for this warehouse.
         Get the default localisation types and set each reference from the
@@ -406,6 +406,10 @@ class Warehouse(Address, DocumentItem, GeoItem, CompleteIdentifierItem,
                 return None, _("No value")
             return
 
+        import_object = None
+        if context and "import_object" in context:
+            import_object = context["import_object"]
+
         TMP_SEMI_COLON = "|#|#|"
         value = value.replace("\\;", TMP_SEMI_COLON)  # manage ";" used by a ref
 
@@ -413,23 +417,27 @@ class Warehouse(Address, DocumentItem, GeoItem, CompleteIdentifierItem,
 
         divisions = list(WarehouseDivisionLink.objects.filter(
             warehouse=self).order_by('order'))
-        if len(values) > len(divisions):
-            if return_errors:
-                return str(_("{} values for only {} default divisions set for "
-                             "warehouse {}")).format(
-                    len(values), len(divisions), self.name)
-            return
 
         parent = None
         for idx, value in enumerate(values):
+            if idx >= len(divisions):
+                if return_errors:
+                    return str(
+                        _("{} values for only {} default divisions set for "
+                          "warehouse {}")).format(
+                        len(values), len(divisions), self.name)
+                return
             value = value.replace(TMP_SEMI_COLON, ";").strip()
             if not value or value == "-":
                 continue
-            parent, __ = Container.objects.get_or_create(
+            parent, created = Container.objects.get_or_create(
                 location=self,
                 reference=value,
                 container_type_id=divisions[idx].container_type_id,
                 parent=parent)
+            if created and import_object:
+                parent.imports.add(import_object)
+
 
     @property
     def short_label(self):