Allow plugins to specify their supported modes

Author: radarhere

Tests/test_image.py

@@ -2,6 +2,8 @@
 
 from PIL import Image, TiffImagePlugin
 from PIL._util import py3
+
+from io import BytesIO
 import os
 import sys
 import shutil
@@ -80,9 +82,7 @@ def test_width_height(self):
 
     def test_invalid_image(self):
         if py3:
-            import io
-
-            im = io.BytesIO(b"")
+            im = BytesIO(b"")
         else:
             import StringIO
 
@@ -359,14 +359,54 @@ def test_registered_extensions(self):
         for ext in [".cur", ".icns", ".tif", ".tiff"]:
             self.assertIn(ext, extensions)
 
-    def test_no_convert_mode(self):
-        self.assertTrue(not hasattr(TiffImagePlugin, '_convert_mode'))
+    def test_supported_modes(self):
+        for format in Image.MIME.keys():
+            try:
+                save_handler = Image.SAVE[format]
+            except KeyError:
+                continue
+            plugin = sys.modules[save_handler.__module__]
+            if not hasattr(plugin, '_supported_modes'):
+                continue
+
+            # Check that the supported modes list is accurate
+            supported_modes = plugin._supported_modes()
+            for mode in ['1', 'L', 'P', 'RGB', 'RGBA', 'CMYK', 'YCbCr', 'LAB',
+                         'HSV', 'I', 'F', 'LA', 'La', 'RGBX', 'RGBa']:
+                out = BytesIO()
+                im = Image.new(mode, (100, 100))
+                if mode in supported_modes:
+                    im.save(out, format)
+                else:
+                    self.assertRaises(Exception, im.save, out, format)
+
+    def test_no_supported_modes_method(self):
+        self.assertTrue(not hasattr(TiffImagePlugin, '_supported_modes'))
 
         temp_file = self.tempfile("temp.tiff")
 
         im = hopper()
         im.save(temp_file, convert_mode=True)
 
+    def test_convert_mode(self):
+        for mode, modes in [
+            ['P', []],  # no modes
+            ['P', ['P']]  # same mode
+        ]:
+            im = Image.new(mode, (100, 100))
+            self.assertIsNone(im._convert_mode(modes))
+
+        for mode, modes in [
+            ['P', ['RGB']],
+            ['P', ['L']],  # converting to a non-preferred mode
+            ['LA', ['P']],
+            ['I', ['L']],
+            ['RGB', ['L']],
+            ['RGB', ['CMYK']]
+        ]:
+            im = Image.new(mode, (100, 100))
+            self.assertIsNotNone(im._convert_mode(modes))
+
     def test_effect_mandelbrot(self):
         # Arrange
         size = (512, 512)

src/PIL/GifImagePlugin.py

@@ -843,11 +843,8 @@ def write(self, data):
     return fp.data
 
 
-def _convert_mode(im):
-    return {
-        'LA':'P',
-        'CMYK':'RGB'
-    }.get(im.mode)
+def _supported_modes():
+    return ["RGB", "RGBA", "P", "I", "F", "LA", "L", "1"]
 
 
 # --------------------------------------------------------------------

src/PIL/Image.py

@@ -2040,16 +2040,18 @@ def save(self, fp, format=None, **params):
 
         if format.upper() not in SAVE:
             init()
-        if params.pop('save_all', False):
+        if params.pop("save_all", False):
             save_handler = SAVE_ALL[format.upper()]
         else:
             save_handler = SAVE[format.upper()]
 
-        if params.get('convert_mode'):
+        if params.get("convert_mode"):
             plugin = sys.modules[save_handler.__module__]
-            converted_im = self._convert_mode(plugin, params)
-            if converted_im:
-                return converted_im.save(fp, format, **params)
+            if hasattr(plugin, "_supported_modes"):
+                modes = plugin._supported_modes()
+                converted_im = self._convert_mode(modes, params)
+                if converted_im:
+                    return converted_im.save(fp, format, **params)
 
         self.encoderinfo = params
         self.encoderconfig = ()
@@ -2069,32 +2071,57 @@ def save(self, fp, format=None, **params):
             if open_fp:
                 fp.close()
 
-    def _convert_mode(self, plugin, params):
-        if not hasattr(plugin, '_convert_mode'):
+    def _convert_mode(self, modes, params={}):
+        if not modes or self.mode in modes:
             return
-        new_mode = plugin._convert_mode(self)
-        if self.mode == 'LA' and new_mode == 'P':
-            alpha = self.getchannel('A')
+        if self.mode == "P":
+            preferred_modes = []
+            if "A" in self.im.getpalettemode():
+                preferred_modes.append("RGBA")
+            preferred_modes.append("RGB")
+        else:
+            preferred_modes = {
+                "CMYK": ["RGB"],
+                "RGB": ["CMYK"],
+                "RGBX": ["RGB"],
+                "RGBa": ["RGBA", "RGB"],
+                "RGBA": ["RGB"],
+                "LA": ["RGBA", "P", "L"],
+                "La": ["LA", "L"],
+                "L": ["RGB"],
+                "F": ["I"],
+                "I": ["L", "RGB"],
+                "1": ["L"],
+                "YCbCr": ["RGB"],
+                "LAB": ["RGB"],
+                "HSV": ["RGB"],
+            }.get(self.mode, [])
+        for new_mode in preferred_modes:
+            if new_mode in modes:
+                break
+        else:
+            new_mode = modes[0]
+        if self.mode == "LA" and new_mode == "P":
+            alpha = self.getchannel("A")
             # Convert the image into P mode but only use 255 colors
             # in the palette out of 256.
-            im = self.convert('L') \
-                .convert('P', palette=ADAPTIVE, colors=255)
+            im = self.convert("L").convert("P", palette=ADAPTIVE, colors=255)
             # Set all pixel values below 128 to 255, and the rest to 0.
             mask = eval(alpha, lambda px: 255 if px < 128 else 0)
             # Paste the color of index 255 and use alpha as a mask.
             im.paste(255, mask)
             # The transparency index is 255.
-            im.info['transparency'] = 255
+            im.info["transparency"] = 255
             return im
 
-        elif self.mode == 'I':
-            im = self.point([i//256 for i in range(65536)], 'L')
-            return im.convert(new_mode) if new_mode != 'L' else im
+        elif self.mode == "I":
+            im = self.point([i // 256 for i in range(65536)], "L")
+            return im.convert(new_mode) if new_mode != "L" else im
 
-        elif self.mode in ('RGBA', 'LA') and new_mode in ('RGB', 'L'):
-            fill_color = params.get('fill_color', 'white')
+        elif self.mode in ("RGBA", "LA") and new_mode in ("RGB", "L"):
+            fill_color = params.get("fill_color", "white")
             background = new(new_mode, self.size, fill_color)
-            background.paste(self, self.getchannel('A'))
+            background.paste(self, self.getchannel("A"))
             return background
 
         elif new_mode:

src/PIL/JpegImagePlugin.py

@@ -815,15 +815,8 @@ def jpeg_factory(fp=None, filename=None):
     return im
 
 
-def _convert_mode(im):
-    mode = im.mode
-    if mode == 'P':
-        return 'RGBA' if 'A' in im.im.getpalettemode() else 'RGB'
-    return {
-        'RGBA':'RGB',
-        'LA':'L',
-        'I':'L'
-    }.get(mode)
+def _supported_modes():
+    return ["RGB", "CMYK", "YCbCr", "RGBX", "L", "1"]
 
 
 # ---------------------------------------------------------------------

src/PIL/PngImagePlugin.py

@@ -946,10 +946,8 @@ def append(fp, cid, *data):
     return fp.data
 
 
-def _convert_mode(im):
-    return {
-        'CMYK':'RGB'
-    }.get(im.mode)
+def _supported_modes():
+    return ["RGB", "RGBA", "P", "I", "LA", "L", "1"]
 
 
 # --------------------------------------------------------------------

src/PIL/WebPImagePlugin.py

@@ -350,17 +350,22 @@ def _save(im, fp, filename):
     fp.write(data)
 
 
-def _convert_mode(im):
-    mode = im.mode
-    if mode == 'P':
-        return 'RGBA' if 'A' in im.im.getpalettemode() else 'RGB'
-    return {
-        # Pillow doesn't support L modes for webp for now.
-        'L':'RGB',
-        'LA':'RGBA',
-        'I':'RGB',
-        'CMYK':'RGB'
-    }.get(mode)
+def _supported_modes():
+    return [
+        "RGB",
+        "RGBA",
+        "RGBa",
+        "RGBX",
+        "CMYK",
+        "YCbCr",
+        "HSV",
+        "I",
+        "F",
+        "P",
+        "LA",
+        "L",
+        "1",
+    ]
 
 
 Image.register_open(WebPImageFile.format, WebPImageFile, _accept)