From a6b95930bbda3b6cbd3120dcd0316a0598d6a2ee Mon Sep 17 00:00:00 2001
From: Matt Westcott <matt@west.co.tt>
Date: Tue, 22 Sep 2015 11:58:36 +0100
Subject: [PATCH] Make Document a collection member model

---
 wagtail/wagtailcore/models.py                 | 20 +++++++++++++++++
 .../migrations/0005_document_collection.py    | 22 +++++++++++++++++++
 wagtail/wagtaildocs/models.py                 |  3 ++-
 3 files changed, 44 insertions(+), 1 deletion(-)
 create mode 100644 wagtail/wagtaildocs/migrations/0005_document_collection.py

diff --git a/wagtail/wagtailcore/models.py b/wagtail/wagtailcore/models.py
index 46b9e316d3..090b4b05ca 100644
--- a/wagtail/wagtailcore/models.py
+++ b/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
diff --git a/wagtail/wagtaildocs/migrations/0005_document_collection.py b/wagtail/wagtaildocs/migrations/0005_document_collection.py
new file mode 100644
index 0000000000..1886672fb2
--- /dev/null
+++ b/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,
+        ),
+    ]
diff --git a/wagtail/wagtaildocs/models.py b/wagtail/wagtaildocs/models.py
index 340d105194..b42df33e2d 100644
--- a/wagtail/wagtaildocs/models.py
+++ b/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)