diff --git a/dbfilestorage/models.py b/dbfilestorage/models.py index 979a395..f318795 100644 --- a/dbfilestorage/models.py +++ b/dbfilestorage/models.py @@ -14,4 +14,12 @@ class DBFile(models.Model): b64 = models.TextField() def __unicode__(self): - return unicode(self.name) + return u"{name} <{content_type}>".format( + name=self.name, content_type=self.content_type) + + def save(self, **kwargs): + if self.content_type is None: + # If content type guessing fails, + # use octet stream as a major fallback + self.content_type = "application/octet-stream" + super(DBFile, self).save(**kwargs) diff --git a/dbfilestorage/storage.py b/dbfilestorage/storage.py index f6f30d4..55d8fce 100644 --- a/dbfilestorage/storage.py +++ b/dbfilestorage/storage.py @@ -40,26 +40,14 @@ class DBFileStorage(Storage): :return str: the name(md5) to look up the file by. """ - ct = None if hasattr(content.file, "read"): read_data = content.file.read() else: read_data = content.file.encode('utf8') b64 = read_data.encode('base64') - # Try to get the real content_type if applicable. - try: - if hasattr(content, 'content_type'): - ct = content.content_type - elif hasattr(content.file, 'content_type'): - ct = content.file.content_type - except AttributeError: - pass - - # USE mimetypes.guess_type as a fallback. - if ct is None: - # https://docs.python.org/2/library/mimetypes.html - ct = mimetypes.guess_type(name)[0] + # USE mimetypes.guess_type as an attempt at getting the content type. + ct = mimetypes.guess_type(name)[0] # After we get the mimetype by name potentially, mangle it. name = hashlib.md5(read_data).hexdigest() @@ -81,7 +69,9 @@ class DBFileStorage(Storage): def delete(self, name): assert name, "The name argument is not allowed to be empty." - DBFile.objects.get(pk=name).delete() + # name is the Pk, so it will be unique, deleting on the QS so + # that it fails silently if the file doesn't exist. + DBFile.objects.filter(pk=name).delete() def exists(self, name): return DBFile.objects.filter(pk=name).exists() diff --git a/tests/tests.py b/tests/tests.py index e0bb9ca..7a08051 100644 --- a/tests/tests.py +++ b/tests/tests.py @@ -1,9 +1,11 @@ +# -*- coding: utf-8 -*- import hashlib import os from dbfilestorage.models import DBFile from django.contrib.auth.models import User +from django.core.files.base import ContentFile from django.core.files.storage import default_storage from django.core.urlresolvers import reverse from django.test import TestCase, Client @@ -30,6 +32,15 @@ class DBFileTest(TestCase): """ Test that the file storage uploads and puts in DB Properly """ self.assertTrue(DBFile.objects.filter(name=self.md5).exists()) + def test_content_file(self): + """ Test that this code works with ContentFile as well """ + content_file = ContentFile(u"ΑΔΔGΕΝΕ") + content_file_md5 = hashlib.md5(u"ΑΔΔGΕΝΕ".encode('utf8')).hexdigest() + default_storage.save("unicode", content_file) + unicode_file = DBFile.objects.get(name=content_file_md5) + self.assertEqual(unicode(unicode_file), + "{} ".format(content_file_md5)) + def test_no_duplicate_upload(self): """ Test that it won't make a new file if it already exists """ # uploads once in setup already @@ -40,8 +51,7 @@ class DBFileTest(TestCase): """ Test that the DB entry matches what is expected from the file """ with open(self.filepath, 'rb') as f: dbf = DBFile.objects.get(name=self.md5) - self.assertEqual(dbf.b64.decode("base64"), - f.read()) + self.assertEqual(dbf.b64.decode("base64"), f.read()) self.assertEqual(dbf.content_type, 'image/jpeg') def test_open(self): @@ -59,6 +69,10 @@ class DBFileTest(TestCase): self.assertTrue(DBFile.objects.filter(name=self.md5).exists()) default_storage.delete(self.md5) self.assertFalse(DBFile.objects.filter(name=self.md5).exists()) + # Also test that calling delete on something that doesn't exist, + # errors silently + self.assertFalse(DBFile.objects.filter(name="Nothing").exists()) + default_storage.delete("Nothing") def test_path(self): """ Test the path is just the md5 name """