Merge pull request #20 from torchbox/wagtail-2-8-compatiblity

Update get_document_model usage for Wagtail >= 2.8 compatibility
torchbox · Aug 4, 2020 · 6c829bd · 6c829bd
2 parents c7359ef + d7770d3
commit 6c829bd
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 3 deletions.
diff --git a/wagtail_storages/factories.py b/wagtail_storages/factories.py
@@ -1,9 +1,14 @@
+from wagtail import VERSION as WAGTAIL_VERSION
 from wagtail.core.models import Collection, CollectionViewRestriction
-from wagtail.documents.models import get_document_model
 
 import factory
 import factory.django
 
+if WAGTAIL_VERSION < (2, 8):
+    from wagtail.documents.models import get_document_model
+else:
+    from wagtail.documents import get_document_model
+
 Document = get_document_model()
 
 

diff --git a/wagtail_storages/signal_handlers.py b/wagtail_storages/signal_handlers.py
@@ -3,8 +3,13 @@
 
 from django.db.models.signals import post_save, pre_delete
 
+from wagtail import VERSION as WAGTAIL_VERSION
 from wagtail.core.models import Collection
-from wagtail.documents.models import get_document_model
+
+if WAGTAIL_VERSION < (2, 8):
+    from wagtail.documents.models import get_document_model
+else:
+    from wagtail.documents import get_document_model
 
 from wagtail_storages.utils import (
     is_s3_boto3_storage_used,

diff --git a/wagtail_storages/utils.py b/wagtail_storages/utils.py
@@ -3,9 +3,14 @@
 from django.conf import settings
 from django.core.files.storage import get_storage_class
 
+from wagtail import VERSION as WAGTAIL_VERSION
 from wagtail.contrib.frontend_cache.utils import PurgeBatch
 from wagtail.core.models import Site
-from wagtail.documents.models import get_document_model
+
+if WAGTAIL_VERSION < (2, 8):
+    from wagtail.documents.models import get_document_model
+else:
+    from wagtail.documents import get_document_model
 
 import storages.backends.s3boto3