Make Document a collection member model

wagtail/wagtailcore/models.py

@@ -1742,3 +1742,23 @@ class Collection(MP_Node):
     class Meta:
         verbose_name = _('collection')
         verbose_name_plural = _('collections')
+
+
+def get_root_collection_id():
+    return Collection.get_first_root_node().id
+
+
+class CollectionMember(models.Model):
+    """
+    Base class for models that are categorised into collections
+    """
+    collection = models.ForeignKey(
+        Collection,
+        default=get_root_collection_id,
+        verbose_name=_('collection'),
+        related_name='+',
+        on_delete=models.CASCADE
+    )
+
+    class Meta:
+        abstract = True

wagtail/wagtaildocs/migrations/0005_document_collection.py

@@ -0,0 +1,22 @@
+# -*- coding: utf-8 -*-
+from __future__ import unicode_literals
+
+from django.db import models, migrations
+import wagtail.wagtailcore.models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('wagtailcore', '0025_collection_initial_data'),
+        ('wagtaildocs', '0004_capitalizeverbose'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='document',
+            name='collection',
+            field=models.ForeignKey(related_name='+', to='wagtailcore.Collection', verbose_name='collection', default=wagtail.wagtailcore.models.get_root_collection_id),
+            preserve_default=True,
+        ),
+    ]

wagtail/wagtaildocs/models.py

@@ -16,6 +16,7 @@ from django.utils.encoding import python_2_unicode_compatible
 
 from wagtail.wagtailadmin.taggable import TagSearchable
 from wagtail.wagtailadmin.utils import get_object_usage
+from wagtail.wagtailcore.models import CollectionMember
 from wagtail.wagtailsearch import index
 from wagtail.wagtailsearch.queryset import SearchableQuerySetMixin
 
@@ -25,7 +26,7 @@ class DocumentQuerySet(SearchableQuerySetMixin, models.QuerySet):
 
 
 @python_2_unicode_compatible
-class AbstractDocument(models.Model, TagSearchable):
+class AbstractDocument(CollectionMember, 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)