diff --git a/wagtail/wagtaildocs/models.py b/wagtail/wagtaildocs/models.py index afc9142d0..2149320f7 100644 --- a/wagtail/wagtaildocs/models.py +++ b/wagtail/wagtaildocs/models.py @@ -37,6 +37,14 @@ class Document(models.Model, TagSearchable): def filename(self): return os.path.basename(self.file.name) + @property + def file_extension(self): + parts = self.filename.split('.') + if len(parts) > 1: + return parts[-1] + else: + return '' + @property def url(self): return reverse('wagtaildocs_serve', args=[self.id, self.filename]) diff --git a/wagtail/wagtaildocs/tests.py b/wagtail/wagtaildocs/tests.py index d59b5ac25..d068f9a27 100644 --- a/wagtail/wagtaildocs/tests.py +++ b/wagtail/wagtaildocs/tests.py @@ -3,6 +3,7 @@ from wagtail.wagtaildocs import models from wagtail.tests.utils import login from django.contrib.auth.models import User, Group, Permission from django.core.urlresolvers import reverse +from django.core.files.base import ContentFile # TODO: Test serve view @@ -150,3 +151,24 @@ class TestDocumentChooserUploadView(TestCase): def test_status_code(self): self.assertEqual(self.get().status_code, 200) + + +class TestDocumentFilenameProperties(TestCase): + def setUp(self): + self.document = models.Document(title="Test document") + self.document.file.save('example.doc', ContentFile("A boring example document")) + + self.extensionless_document = models.Document(title="Test document") + self.extensionless_document.file.save('example', ContentFile("A boring example document")) + + def test_filename(self): + self.assertEqual('example.doc', self.document.filename) + self.assertEqual('example', self.extensionless_document.filename) + + def test_file_extension(self): + self.assertEqual('doc', self.document.file_extension) + self.assertEqual('', self.extensionless_document.file_extension) + + def tearDown(self): + self.document.delete() + self.extensionless_document.delete()