Container: put_document_by_external_id, put_document_by_reference

1 files changed, 21 insertions, 5 deletions

diff --git a/archaeological_warehouse/models.py b/archaeological_warehouse/models.py
index cb144270b..e094436e1 100644
--- a/archaeological_warehouse/models.py
+++ b/archaeological_warehouse/models.py
@@ -1000,18 +1000,34 @@ class Container(DocumentItem, Merge, LightHistorizedItem,
         except Warehouse.DoesNotExist:
             return
 
-    @post_importer_action
-    def put_document_by_complete_identifier(self, context, value):
+    def put_document_by_key(self, value, key):
         Document = apps.get_model("ishtar_common", "Document")
         try:
-            doc = Document.objects.get(complete_identifier=value)
+            doc = Document.objects.get(**{key: value})
         except Document.DoesNotExist:
-            raise ImporterError(str(_("Document with complete id: {} does not "
-                                      "exists")).format(value))
+            raise ImporterError(str(_("Document with {}: {} does not "
+                                      "exists")).format(key, value))
+        except Document.MultipleObjectsReturned:
+            raise ImporterError(
+                str(_("Multiple document with {}: {}")).format(key, value))
         doc.container_id = self.pk
         doc.container_ref_id = self.pk
         doc.skip_history_when_saving = True
         doc.save()
+
+    @post_importer_action
+    def put_document_by_external_id(self, context, value):
+        self.put_document_by_key(value, "external_id")
+    put_document_by_external_id.post_save = True
+
+    @post_importer_action
+    def put_document_by_reference(self, context, value):
+        self.put_document_by_key(value, "reference")
+    put_document_by_reference.post_save = True
+
+    @post_importer_action
+    def put_document_by_complete_identifier(self, context, value):
+        self.put_document_by_key(value, "complete_identifier")
     put_document_by_complete_identifier.post_save = True
 
     def _generate_cached_division(self):