Better handle types when running optimisers (#148)

The main change here is adding an explicit type error, rather than having a variable not be defined.

Secondly, the type handling itself gets simpler and slightly more optimised.

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(

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)

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,42 +148,44 @@ 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)
 
             if hasattr(image_file, "seek"):
                 # 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)