From 27829f6090c7c3a3c81846868801a9f0993e8582 Mon Sep 17 00:00:00 2001
From: Alex Gleason <alex@alexgleason.me>
Date: Wed, 16 Dec 2015 21:09:35 -0600
Subject: [PATCH] Allow overriding document model

Added missed import
---
 wagtail/wagtaildocs/admin.py           | 10 +++++++-
 wagtail/wagtaildocs/blocks.py          |  4 ++--
 wagtail/wagtaildocs/forms.py           | 18 +++++++--------
 wagtail/wagtaildocs/models.py          | 32 +++++++++++++++++++++++++-
 wagtail/wagtaildocs/rich_text.py       |  3 ++-
 wagtail/wagtaildocs/views/chooser.py   | 12 +++++++---
 wagtail/wagtaildocs/views/documents.py | 14 +++++++++--
 wagtail/wagtaildocs/views/serve.py     |  3 ++-
 wagtail/wagtaildocs/wagtail_hooks.py   |  4 ++--
 wagtail/wagtaildocs/widgets.py         |  8 +++++--
 10 files changed, 83 insertions(+), 25 deletions(-)

diff --git a/wagtail/wagtaildocs/admin.py b/wagtail/wagtaildocs/admin.py
index 8c295cd6f..83af3e964 100644
--- a/wagtail/wagtaildocs/admin.py
+++ b/wagtail/wagtaildocs/admin.py
@@ -1,5 +1,13 @@
 from django.contrib import admin
+from django.conf import settings
 
 from wagtail.wagtaildocs.models import Document
 
-admin.site.register(Document)
+
+if hasattr(settings, 'WAGTAILDOCS_DOCUMENT_MODEL') and settings.WAGTAILDOCS_DOCUMENT_MODEL != 'wagtaildocs.Document':
+    # This installation provides its own custom document class;
+    # to avoid confusion, we won't expose the unused wagtaildocs.Document class
+    # in the admin.
+    pass
+else:
+    admin.site.register(Document)
diff --git a/wagtail/wagtaildocs/blocks.py b/wagtail/wagtaildocs/blocks.py
index 8db11fd8a..bea125754 100644
--- a/wagtail/wagtaildocs/blocks.py
+++ b/wagtail/wagtaildocs/blocks.py
@@ -9,8 +9,8 @@ from wagtail.wagtailcore.blocks import ChooserBlock
 class DocumentChooserBlock(ChooserBlock):
     @cached_property
     def target_model(self):
-        from wagtail.wagtaildocs.models import Document
-        return Document
+        from wagtail.wagtaildocs.models import get_document_model
+        return get_document_model()
 
     @cached_property
     def widget(self):
diff --git a/wagtail/wagtaildocs/forms.py b/wagtail/wagtaildocs/forms.py
index 90d4a9880..60d100be0 100644
--- a/wagtail/wagtaildocs/forms.py
+++ b/wagtail/wagtaildocs/forms.py
@@ -1,16 +1,14 @@
 from django import forms
+from django.forms.models import modelform_factory
 
 from wagtail.wagtailadmin import widgets
-from wagtail.wagtaildocs.models import Document
 
 
-class DocumentForm(forms.ModelForm):
-    required_css_class = "required"
-
-    class Meta:
-        model = Document
-        fields = ('title', 'file', 'tags')
-        widgets = {
-            'file': forms.FileInput(),
+def get_document_form(model):
+    return modelform_factory(
+        model,
+        fields=model.admin_form_fields,
+        widgets={
             'tags': widgets.AdminTagWidget,
-        }
+            'file': forms.FileInput()
+        })
diff --git a/wagtail/wagtaildocs/models.py b/wagtail/wagtaildocs/models.py
index dfc109e76..1b4cdf119 100644
--- a/wagtail/wagtaildocs/models.py
+++ b/wagtail/wagtaildocs/models.py
@@ -9,6 +9,7 @@ from django.db.models.signals import pre_delete
 from django.dispatch.dispatcher import receiver
 from django.dispatch import Signal
 from django.core.urlresolvers import reverse
+from django.core.exceptions import ImproperlyConfigured
 from django.conf import settings
 from django.utils.translation import ugettext_lazy as _
 from django.utils.encoding import python_2_unicode_compatible
@@ -24,7 +25,7 @@ class DocumentQuerySet(SearchableQuerySetMixin, models.QuerySet):
 
 
 @python_2_unicode_compatible
-class Document(models.Model, TagSearchable):
+class AbstractDocument(models.Model, TagSearchable):
     title = models.CharField(max_length=255, verbose_name=_('title'))
     file = models.FileField(upload_to='documents', verbose_name=_('file'))
     created_at = models.DateTimeField(verbose_name=_('created at'), auto_now_add=True)
@@ -78,9 +79,38 @@ class Document(models.Model, TagSearchable):
             return False
 
     class Meta:
+        abstract = True
         verbose_name = _('document')
 
 
+class Document(AbstractDocument):
+    admin_form_fields = (
+        'title',
+        'file',
+        'tags'
+    )
+
+
+def get_document_model():
+    from django.conf import settings
+    from django.apps import apps
+
+    try:
+        app_label, model_name = settings.WAGTAILDOCS_DOCUMENT_MODEL.split('.')
+    except AttributeError:
+        return Document
+    except ValueError:
+        raise ImproperlyConfigured("WAGTAILDOCS_DOCUMENT_MODEL must be of the form 'app_label.model_name'")
+
+    document_model = apps.get_model(app_label, model_name)
+    if document_model is None:
+        raise ImproperlyConfigured(
+            "WAGTAILDOCS_DOCUMENT_MODEL refers to model '%s' that has not been installed" %
+            settings.WAGTAILDOCS_DOCUMENT_MODEL
+        )
+    return document_model
+
+
 # Receive the pre_delete signal and delete the file associated with the model instance.
 @receiver(pre_delete, sender=Document)
 def document_delete(sender, instance, **kwargs):
diff --git a/wagtail/wagtaildocs/rich_text.py b/wagtail/wagtaildocs/rich_text.py
index 44f42a2e9..3c5013f71 100644
--- a/wagtail/wagtaildocs/rich_text.py
+++ b/wagtail/wagtaildocs/rich_text.py
@@ -1,6 +1,6 @@
 from django.utils.html import escape
 
-from wagtail.wagtaildocs.models import Document
+from wagtail.wagtaildocs.models import get_document_model
 
 
 class DocumentLinkHandler(object):
@@ -10,6 +10,7 @@ class DocumentLinkHandler(object):
 
     @staticmethod
     def expand_db_attributes(attrs, for_editor):
+        Document = get_document_model()
         try:
             doc = Document.objects.get(id=attrs['id'])
 
diff --git a/wagtail/wagtaildocs/views/chooser.py b/wagtail/wagtaildocs/views/chooser.py
index 421e7a7ec..488b82802 100644
--- a/wagtail/wagtaildocs/views/chooser.py
+++ b/wagtail/wagtaildocs/views/chooser.py
@@ -9,8 +9,8 @@ from wagtail.wagtailadmin.forms import SearchForm
 from wagtail.wagtailadmin.utils import permission_required
 from wagtail.wagtailsearch.backends import get_search_backends
 
-from wagtail.wagtaildocs.models import Document
-from wagtail.wagtaildocs.forms import DocumentForm
+from wagtail.wagtaildocs.models import get_document_model
+from wagtail.wagtaildocs.forms import get_document_form
 
 
 def get_document_json(document):
@@ -27,7 +27,10 @@ def get_document_json(document):
 
 
 def chooser(request):
+    Document = get_document_model()
+
     if request.user.has_perm('wagtaildocs.add_document'):
+        DocumentForm = get_document_form(Document)
         uploadform = DocumentForm()
     else:
         uploadform = None
@@ -70,7 +73,7 @@ def chooser(request):
 
 
 def document_chosen(request, document_id):
-    document = get_object_or_404(Document, id=document_id)
+    document = get_object_or_404(get_document_model(), id=document_id)
 
     return render_modal_workflow(
         request, None, 'wagtaildocs/chooser/document_chosen.js',
@@ -80,6 +83,9 @@ def document_chosen(request, document_id):
 
 @permission_required('wagtaildocs.add_document')
 def chooser_upload(request):
+    Document = get_document_model()
+    DocumentForm = get_document_form(Document)
+
     if request.POST:
         document = Document(uploaded_by_user=request.user)
         form = DocumentForm(request.POST, request.FILES, instance=document)
diff --git a/wagtail/wagtaildocs/views/documents.py b/wagtail/wagtaildocs/views/documents.py
index b536b110b..cf8b62117 100644
--- a/wagtail/wagtaildocs/views/documents.py
+++ b/wagtail/wagtaildocs/views/documents.py
@@ -10,13 +10,15 @@ from wagtail.wagtailadmin.utils import permission_required, any_permission_requi
 from wagtail.wagtailsearch.backends import get_search_backends
 from wagtail.wagtailadmin import messages
 
-from wagtail.wagtaildocs.models import Document
-from wagtail.wagtaildocs.forms import DocumentForm
+from wagtail.wagtaildocs.models import get_document_model
+from wagtail.wagtaildocs.forms import get_document_form
 
 
 @any_permission_required('wagtaildocs.add_document', 'wagtaildocs.change_document')
 @vary_on_headers('X-Requested-With')
 def index(request):
+    Document = get_document_model()
+
     # Get documents
     documents = Document.objects.all()
 
@@ -67,6 +69,9 @@ def index(request):
 
 @permission_required('wagtaildocs.add_document')
 def add(request):
+    Document = get_document_model()
+    DocumentForm = get_document_form(Document)
+
     if request.POST:
         doc = Document(uploaded_by_user=request.user)
         form = DocumentForm(request.POST, request.FILES, instance=doc)
@@ -92,6 +97,9 @@ def add(request):
 
 
 def edit(request, document_id):
+    Document = get_document_model()
+    DocumentForm = get_document_form(Document)
+
     doc = get_object_or_404(Document, id=document_id)
 
     if not doc.is_editable_by_user(request.user):
@@ -146,6 +154,7 @@ def edit(request, document_id):
 
 
 def delete(request, document_id):
+    Document = get_document_model()
     doc = get_object_or_404(Document, id=document_id)
 
     if not doc.is_editable_by_user(request.user):
@@ -162,6 +171,7 @@ def delete(request, document_id):
 
 
 def usage(request, document_id):
+    Document = get_document_model()
     doc = get_object_or_404(Document, id=document_id)
 
     paginator, used_by = paginate(request, doc.get_usage())
diff --git a/wagtail/wagtaildocs/views/serve.py b/wagtail/wagtaildocs/views/serve.py
index a712b1a8e..bf3842ef1 100644
--- a/wagtail/wagtaildocs/views/serve.py
+++ b/wagtail/wagtaildocs/views/serve.py
@@ -8,10 +8,11 @@ from wsgiref.util import FileWrapper
 from wagtail.utils.sendfile import sendfile
 from wagtail.utils import sendfile_streaming_backend
 
-from wagtail.wagtaildocs.models import Document, document_served
+from wagtail.wagtaildocs.models import get_document_model, document_served
 
 
 def serve(request, document_id, document_filename):
+    Document = get_document_model()
     doc = get_object_or_404(Document, id=document_id)
 
     # Send document_served signal
diff --git a/wagtail/wagtaildocs/wagtail_hooks.py b/wagtail/wagtaildocs/wagtail_hooks.py
index 9a637d1c2..957c7b87a 100644
--- a/wagtail/wagtaildocs/wagtail_hooks.py
+++ b/wagtail/wagtaildocs/wagtail_hooks.py
@@ -11,7 +11,7 @@ from wagtail.wagtailadmin.site_summary import SummaryItem
 from wagtail.wagtailadmin.search import SearchArea
 
 from wagtail.wagtaildocs import admin_urls
-from wagtail.wagtaildocs.models import Document
+from wagtail.wagtaildocs.models import get_document_model
 from wagtail.wagtaildocs.rich_text import DocumentLinkHandler
 
 
@@ -76,7 +76,7 @@ class DocumentsSummaryItem(SummaryItem):
 
     def get_context(self):
         return {
-            'total_docs': Document.objects.count(),
+            'total_docs': get_document_model().objects.count(),
         }
 
 
diff --git a/wagtail/wagtaildocs/widgets.py b/wagtail/wagtaildocs/widgets.py
index a6ca76be0..0b94f0c77 100644
--- a/wagtail/wagtaildocs/widgets.py
+++ b/wagtail/wagtaildocs/widgets.py
@@ -6,7 +6,7 @@ from django.template.loader import render_to_string
 from django.utils.translation import ugettext_lazy as _
 
 from wagtail.wagtailadmin.widgets import AdminChooser
-from wagtail.wagtaildocs.models import Document
+from wagtail.wagtaildocs.models import get_document_model
 
 
 class AdminDocumentChooser(AdminChooser):
@@ -14,8 +14,12 @@ class AdminDocumentChooser(AdminChooser):
     choose_another_text = _('Choose another document')
     link_to_chosen_text = _('Edit this document')
 
+    def __init__(self, **kwargs):
+        super(AdminDocumentChooser, self).__init__(**kwargs)
+        self.document_model = get_document_model()
+
     def render_html(self, name, value, attrs):
-        instance, value = self.get_instance_and_id(Document, value)
+        instance, value = self.get_instance_and_id(self.document_model, value)
         original_field_html = super(AdminDocumentChooser, self).render_html(name, value, attrs)
 
         return render_to_string("wagtaildocs/widgets/document_chooser.html", {