Skip to content
Permalink
Browse files Browse the repository at this point in the history
Don't load temporary uploaded files into memory
These can be quite large. Instead pass them straight to Willow to have it read just as much as it needs to
  • Loading branch information
RealOrangeOne authored and gasman committed Apr 3, 2023
1 parent 3c0c646 commit cfa11bb
Show file tree
Hide file tree
Showing 2 changed files with 43 additions and 5 deletions.
8 changes: 4 additions & 4 deletions wagtail/images/fields.py
Expand Up @@ -146,11 +146,11 @@ def to_python(self, data):
if f is None:
return None

# We need to get a file object for Pillow. When we get a path, we need to open
# the file first. And we have to read the data into memory to pass to Willow.
# Get the file content ready for Willow
if hasattr(data, "temporary_file_path"):
with open(data.temporary_file_path(), "rb") as fh:
file = BytesIO(fh.read())
# Django's `TemporaryUploadedFile` is enough of a file to satisfy Willow
# Willow doesn't support opening images by path https://github.com/wagtail/Willow/issues/108
file = data
else:
if hasattr(data, "read"):
file = BytesIO(data.read())
Expand Down
40 changes: 39 additions & 1 deletion wagtail/images/tests/test_admin_views.py
Expand Up @@ -3,7 +3,7 @@
import urllib

from django.contrib.auth.models import Group, Permission
from django.core.files.uploadedfile import SimpleUploadedFile
from django.core.files.uploadedfile import SimpleUploadedFile, TemporaryUploadedFile
from django.template.defaultfilters import filesizeformat
from django.template.loader import render_to_string
from django.test import RequestFactory, TestCase, override_settings
Expand Down Expand Up @@ -567,6 +567,44 @@ def test_add_svg(self):
images = Image.objects.filter(title="Test image")
self.assertEqual(images.count(), 1)

def test_add_temporary_uploaded_file(self):
"""
Test that uploading large files (spooled to the filesystem) work as expected
"""
test_image_file = get_test_image_file()
uploaded_file = TemporaryUploadedFile(
"test.png", "image/png", test_image_file.size, "utf-8"
)
uploaded_file.write(test_image_file.file.getvalue())
uploaded_file.seek(0)

response = self.post(
{
"title": "Test image",
"file": uploaded_file,
}
)

# Should redirect back to index
self.assertRedirects(response, reverse("wagtailimages:index"))

# Check that the image was created
images = Image.objects.filter(title="Test image")
self.assertEqual(images.count(), 1)

# Test that size was populated correctly
image = images.first()
self.assertEqual(image.width, 640)
self.assertEqual(image.height, 480)

# Test that the file_size/hash fields were set
self.assertTrue(image.file_size)
self.assertTrue(image.file_hash)

# Test that it was placed in the root collection
root_collection = Collection.get_first_root_node()
self.assertEqual(image.collection, root_collection)

@override_settings(
DEFAULT_FILE_STORAGE="wagtail.test.dummy_external_storage.DummyExternalStorage"
)
Expand Down

0 comments on commit cfa11bb

Please sign in to comment.