diff --git a/tests/test_image.py b/tests/test_image.py index f4e17b6..194c9b4 100644 --- a/tests/test_image.py +++ b/tests/test_image.py @@ -282,6 +282,13 @@ def test_optimize_with_file_path(self, mock_process): self.image.optimize("outfile", "jpeg") mock_process.assert_called_with("outfile") + def test_optimize_with_unrecognised_type(self, mock_process): + with self.assertRaises(TypeError): + self.image.optimize(None, "jpeg") + with self.assertRaises(TypeError): + self.image.optimize(io.StringIO(), "jpeg") + mock_process.assert_not_called() + @mock.patch("willow.image.NamedTemporaryFile") @mock.patch("willow.image.os.unlink") def test_optimize_with_bytes( diff --git a/tests/test_pillow.py b/tests/test_pillow.py index 4c346b9..df7faac 100644 --- a/tests/test_pillow.py +++ b/tests/test_pillow.py @@ -554,11 +554,11 @@ def test_save_as_png(self): @unittest.skipIf(not Gifsicle.check_library(), "gifsicle not installed") def test_save_as_gif(self): with open("tests/images/transparent.gif", "rb") as f: - original_size = f.tell() + original_size = os.fstat(f.fileno()).st_size image = PillowImage.open(GIFImageFile(f)) return_value = image.save_as_gif(io.BytesIO()) - self.assertTrue(original_size < return_value.f.tell()) + self.assertTrue(original_size < return_value.f.seek(0, io.SEEK_END)) with mock.patch("willow.plugins.pillow.PillowImage.optimize") as mock_optimize: image.save_as_gif(io.BytesIO(), apply_optimizers=False) @@ -574,7 +574,7 @@ def test_save_as_webp(self): image = PillowImage.open(WebPImageFile(f)) return_value = image.save_as_gif(io.BytesIO()) - self.assertTrue(original_size < return_value.f.tell()) + self.assertTrue(original_size < return_value.f.seek(0, io.SEEK_END)) with mock.patch("willow.plugins.pillow.PillowImage.optimize") as mock_optimize: image.save_as_webp(io.BytesIO(), apply_optimizers=False) diff --git a/willow/image.py b/willow/image.py index 4b425f4..c3f3f2c 100644 --- a/willow/image.py +++ b/willow/image.py @@ -1,6 +1,7 @@ import os import re from io import BytesIO +from shutil import copyfileobj from tempfile import NamedTemporaryFile, SpooledTemporaryFile from typing import Optional @@ -147,31 +148,32 @@ def optimize(self, image_file, image_format: str): named_file_created = False try: - if isinstance(image_file, SpooledTemporaryFile): - file = image_file._file + if isinstance(image_file, (SpooledTemporaryFile, BytesIO)): with NamedTemporaryFile(delete=False) as named_file: - if hasattr(file, "getvalue"): # e.g. BytesIO - named_file.write(file.getvalue()) - else: # e.g. BufferedRandom - file.seek(0) - named_file.write(file.read()) - file_path = named_file.name - named_file_created = True - elif isinstance(image_file, BytesIO): - with NamedTemporaryFile(delete=False) as named_file: - named_file.write(image_file.getvalue()) + named_file_created = True + + image_file.seek(0) + copyfileobj(image_file, named_file) + file_path = named_file.name - named_file_created = True + elif hasattr(image_file, "name"): file_path = image_file.name + elif isinstance(image_file, str): file_path = image_file + elif isinstance(image_file, bytes): with NamedTemporaryFile(delete=False) as named_file: named_file.write(image_file) file_path = named_file.name named_file_created = True + else: + raise TypeError( + f"Cannot optimise {type(image_file)}. It must be a readable object, or a path to a file" + ) + for optimizer in optimizers: optimizer.process(file_path) @@ -179,10 +181,11 @@ def optimize(self, image_file, image_format: str): # rewind and replace the image file with the optimized version image_file.seek(0) with open(file_path, "rb") as f: - image_file.write(f.read()) + copyfileobj(f, image_file) + + if hasattr(image_file, "truncate"): + image_file.truncate() # bring the file size down to the actual image size - if hasattr(image_file, "truncate"): - image_file.truncate() # bring the file size down to the actual image size finally: if named_file_created: os.unlink(file_path)