Skip to content

Instantly share code, notes, and snippets.

@sodastsai
Last active October 15, 2018 10:04
Show Gist options
  • Save sodastsai/5ef3d2c14816d75c5af37158f246947a to your computer and use it in GitHub Desktop.
Save sodastsai/5ef3d2c14816d75c5af37158f246947a to your computer and use it in GitHub Desktop.
Django's storage class which would override files
from django.core.files.storage import FileSystemStorage as DjangoFileSystemStorage
class OverrideStorageMixin(object):
# noinspection PyMethodMayBeStatic,PyUnusedLocal
def get_available_name(self, name, max_length=None):
return name
# noinspection PyUnresolvedReferences
def _save(self, name, content):
# Remove existed file before saving
if self.exists(name):
self.delete(name)
# noinspection PyProtectedMember
return super(OverrideStorageMixin, self)._save(name, content)
class FileSystemStorage(OverrideStorageMixin, DjangoFileSystemStorage):
pass
from io import StringIO
from django.core.files.storage import Storage
from django.test import SimpleTestCase as TestCase
from override_storages import OverrideStorageMixin
# noinspection PyAbstractClass
class InMemoryStorage(Storage):
def __init__(self):
super(InMemoryStorage, self).__init__()
self.file_pool = {}
def delete(self, name):
del self.file_pool[name]
def exists(self, name):
return name in self.file_pool
def _open(self, name, mode="r"):
if mode.startswith("r"):
return StringIO(self.file_pool[name])
else:
raise NotImplementedError("Cannot open file for write")
def _save(self, name, content):
self.file_pool[name] = content.read()
return name
# noinspection PyAbstractClass
class OverrideInMemoryStorage(OverrideStorageMixin, InMemoryStorage):
pass
class InMemoryStorageTests(TestCase):
def setUp(self):
super(InMemoryStorageTests, self).setUp()
self.storage = InMemoryStorage()
def save_file(self, file_name, content_str):
content = StringIO(content_str)
final_file_name = self.storage.save(file_name, content)
content.close()
return final_file_name
def test_normal_readwrite(self):
first_file_name = self.save_file("xd", "1234")
self.assertEqual("xd", first_file_name)
with self.storage.open("xd") as f:
self.assertEqual("1234", f.read())
def test_exists(self):
self.save_file("xd", "1234")
self.assertTrue(self.storage.exists("xd"))
self.assertFalse(self.storage.exists("yo"))
def test_delete(self):
self.save_file("xd", "1234")
self.assertTrue(self.storage.exists("xd"))
self.storage.delete("xd")
self.assertFalse(self.storage.exists("xd"))
class OverrideStorageMixinTests(TestCase):
@staticmethod
def save_file(storage, file_name, content_str):
content = StringIO(content_str)
final_file_name = storage.save(file_name, content)
content.close()
return final_file_name
def test_without_override_mixin(self):
storage = InMemoryStorage()
first_file_name = self.save_file(storage, "xd", "1234")
second_file_name = self.save_file(storage, "xd", "5678")
self.assertEqual("xd", first_file_name)
with storage.open(first_file_name) as f:
self.assertEqual("1234", f.read())
self.assertTrue(second_file_name.startswith("xd_"))
with storage.open(second_file_name) as f:
self.assertEqual("5678", f.read())
def test_with_override_mixin(self):
storage = OverrideInMemoryStorage()
first_file_name = self.save_file(storage, "xd", "1234")
second_file_name = self.save_file(storage, "xd", "5678")
# NOTE: The content of first file has been overridden by the second save operation.
self.assertEqual("xd", first_file_name)
with storage.open(first_file_name) as f:
self.assertEqual("5678", f.read())
self.assertEqual("xd", second_file_name)
with storage.open(second_file_name) as f:
self.assertEqual("5678", f.read())
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment