Add file_size field to Document

CHANGELOG.txt

@@ -9,6 +9,7 @@ Changelog
  * Pillow's image optimisation is now applied when saving PNG images (Dmitry Vasilev)
  * JS / CSS media files can now be associated with Draftail feature definitions (Matt Westcott)
  * The `{% slugurl %}` template tag is now site-aware (Samir Shah)
+ * Added `file_size` field to documents (Karl Hobley)
  * Fix: Handle all exceptions from `Image.get_file_size` (Andrew Plummer)
  * Fix: Fix display of breadcrumbs in ModelAdmin (LB (Ben Johnston))
 

docs/releases/2.2.rst

@@ -18,6 +18,7 @@ Other features
  * Pillow's image optimisation is now applied when saving PNG images (Dmitry Vasilev)
  * JS / CSS media files can now be associated with Draftail feature definitions (Matt Westcott)
  * The ``{% slugurl %}`` template tag is now site-aware (Samir Shah)
+ * Added ``file_size`` field to documents (Karl Hobley)
 
 Bug fixes
 ~~~~~~~~~

wagtail/documents/migrations/0008_document_file_size.py

@@ -0,0 +1,18 @@
+# Generated by Django 2.0.4 on 2018-05-04 17:13
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('wagtaildocs', '0007_merge'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='document',
+            name='file_size',
+            field=models.PositiveIntegerField(editable=False, null=True),
+        ),
+    ]

wagtail/documents/models.py

@@ -33,6 +33,8 @@ class AbstractDocument(CollectionMember, index.Indexed, models.Model):
 
     tags = TaggableManager(help_text=None, blank=True, verbose_name=_('tags'))
 
+    file_size = models.PositiveIntegerField(null=True, editable=False)
+
     objects = DocumentQuerySet.as_manager()
 
     search_fields = CollectionMember.search_fields + [
@@ -44,6 +46,20 @@ class AbstractDocument(CollectionMember, index.Indexed, models.Model):
         index.FilterField('uploaded_by_user'),
     ]
 
+    def get_file_size(self):
+        if self.file_size is None:
+            try:
+                self.file_size = self.file.size
+            except (SystemExit, KeyboardInterrupt):
+                raise
+            except Exception:
+                # File doesn't exist
+                return
+
+            self.save(update_fields=['file_size'])
+
+        return self.file_size
+
     def __str__(self):
         return self.title
 

wagtail/documents/tests/test_admin_views.py

@@ -121,13 +121,16 @@ class TestDocumentAddView(TestCase, WagtailTestUtils):
         self.assertRedirects(response, reverse('wagtaildocs:index'))
 
         # Document should be created, and be placed in the root collection
-        self.assertTrue(models.Document.objects.filter(title="Test document").exists())
+        document = models.Document.objects.get(title="Test document")
         root_collection = Collection.get_first_root_node()
         self.assertEqual(
-            models.Document.objects.get(title="Test document").collection,
+            document.collection,
             root_collection
         )
 
+        # Check that the file_size field was set
+        self.assertTrue(document.file_size)
+
     def test_post_with_collections(self):
         root_collection = Collection.get_first_root_node()
         evil_plans_collection = root_collection.add_child(name="Evil plans")

wagtail/documents/views/chooser.py

@@ -110,6 +110,8 @@ def chooser_upload(request):
         form = DocumentForm(request.POST, request.FILES, instance=document, user=request.user)
 
         if form.is_valid():
+            document.file_size = document.file.size
+
             form.save()
 
             # Reindex the document to make sure all tags are indexed

wagtail/documents/views/documents.py

@@ -1,3 +1,5 @@
+import os
+
 from django.shortcuts import get_object_or_404, redirect, render
 from django.urls import reverse
 from django.utils.translation import ugettext as _
@@ -94,6 +96,8 @@ def add(request):
         doc = Document(uploaded_by_user=request.user)
         form = DocumentForm(request.POST, request.FILES, instance=doc, user=request.user)
         if form.is_valid():
+            doc.file_size = doc.file.size
+
             form.save()
 
             # Reindex the document to make sure all tags are indexed
@@ -129,6 +133,8 @@ def edit(request, document_id):
         if form.is_valid():
             doc = form.save()
             if 'file' in form.changed_data:
+                doc.file_size = doc.file.size
+
                 # if providing a new document file, delete the old one.
                 # NB Doing this via original_file.delete() clears the file field,
                 # which definitely isn't what we want...
@@ -146,26 +152,24 @@ def edit(request, document_id):
     else:
         form = DocumentForm(instance=doc, user=request.user)
 
-    filesize = None
+    try:
+        local_path = doc.file.path
+    except NotImplementedError:
+        # Document is hosted externally (eg, S3)
+        local_path = None
 
-    # Get file size when there is a file associated with the Document object
-    if doc.file:
-        try:
-            filesize = doc.file.size
-        except OSError:
-            # File doesn't exist
-            pass
-
-    if not filesize:
-        messages.error(
-            request,
-            _("The file could not be found. Please change the source or delete the document"),
-            buttons=[messages.button(reverse('wagtaildocs:delete', args=(doc.id,)), _('Delete'))]
-        )
+    if local_path:
+        # Give error if document file doesn't exist
+        if not os.path.isfile(local_path):
+            messages.error(
+                request,
+                _("The file could not be found. Please change the source or delete the document"),
+                buttons=[messages.button(reverse('wagtaildocs:delete', args=(doc.id,)), _('Delete'))]
+            )
 
     return render(request, "wagtaildocs/documents/edit.html", {
         'document': doc,
-        'filesize': filesize,
+        'filesize': doc.get_file_size(),
         'form': form,
         'user_can_delete': permission_policy.user_has_permission_for_instance(
             request.user, 'delete', doc

wagtail/tests/testapp/migrations/0031_customdocument_file_size.py

@@ -0,0 +1,18 @@
+# Generated by Django 2.0.4 on 2018-05-04 17:13
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('tests', '0030_formclassadditionalfieldpage'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='customdocument',
+            name='file_size',
+            field=models.PositiveIntegerField(editable=False, null=True),
+        ),
+    ]