Skip to content

Commit

Permalink
[s3] Fix writing bytearrays
Browse files Browse the repository at this point in the history
  • Loading branch information
jschneier committed Dec 16, 2020
1 parent 01cd42f commit 418b977
Show file tree
Hide file tree
Showing 3 changed files with 24 additions and 4 deletions.
8 changes: 4 additions & 4 deletions storages/backends/s3boto3.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,13 +11,13 @@
from django.core.exceptions import ImproperlyConfigured, SuspiciousOperation
from django.core.files.base import File
from django.utils.deconstruct import deconstructible
from django.utils.encoding import filepath_to_uri, force_bytes
from django.utils.encoding import filepath_to_uri
from django.utils.timezone import is_naive, make_naive

from storages.base import BaseStorage
from storages.utils import (
NonCloseableBufferedReader, check_location, get_available_overwrite_name,
lookup_env, safe_join, setting,
lookup_env, safe_join, setting, to_bytes,
)

try:
Expand Down Expand Up @@ -158,7 +158,7 @@ def write(self, content):
)
if self.buffer_size <= self._buffer_file_size:
self._flush_write_buffer()
bstr = force_bytes(content)
bstr = to_bytes(content)
self._raw_bytes_written += len(bstr)
return super().write(bstr)

Expand Down Expand Up @@ -412,7 +412,7 @@ def _compress_content(self, content):
# For S3 this defeats detection of changes using MD5 sums on gzipped files
# Fixing the mtime at 0.0 at compression time avoids this problem
with GzipFile(mode='wb', fileobj=zbuf, mtime=0.0) as zfile:
zfile.write(force_bytes(content.read()))
zfile.write(to_bytes(content.read()))
zbuf.seek(0)
# Boto 2 returned the InMemoryUploadedFile with the file pointer replaced,
# but Boto 3 seems to have issues with that. No need for fp.name in Boto3
Expand Down
9 changes: 9 additions & 0 deletions storages/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,15 @@
from django.core.exceptions import (
ImproperlyConfigured, SuspiciousFileOperation,
)
from django.utils.encoding import force_bytes


def to_bytes(content):
"""Wrap Django's force_bytes to pass through bytearrays."""
if isinstance(content, bytearray):
return content

return force_bytes(content)


def setting(name, default=None):
Expand Down
11 changes: 11 additions & 0 deletions tests/test_s3boto3.py
Original file line number Diff line number Diff line change
Expand Up @@ -296,6 +296,17 @@ def test_storage_open_write(self):
multipart.complete.assert_called_once_with(
MultipartUpload={'Parts': [{'ETag': '123', 'PartNumber': 1}]})

def test_write_bytearray(self):
"""Test that bytearray write exactly (no extra "bytearray" from stringify)."""
name = "saved_file.bin"
content = bytearray(b"content")
file = self.storage.open(name, "wb")
obj = self.storage.bucket.Object.return_value
# Set the name of the mock object
obj.key = name
bytes_written = file.write(content)
self.assertEqual(len(content), bytes_written)

def test_storage_open_no_write(self):
"""
Test opening file in write mode and closing without writing.
Expand Down

0 comments on commit 418b977

Please sign in to comment.