Improve iPython interoperability

Tests/test_file_jpeg.py

@@ -930,8 +930,7 @@ def test_repr_jpeg(self):
             assert_image_similar(im, repr_jpeg, 17)
 
     def test_repr_jpeg_error_returns_none(self):
-        im = hopper("F")
-
+        im = hopper("BGR;24")
         assert im._repr_jpeg_() is None
 
 

Tests/test_file_png.py

@@ -533,8 +533,7 @@ def test_repr_png(self):
             assert_image_equal(im, repr_png)
 
     def test_repr_png_error_returns_none(self):
-        im = hopper("F")
-
+        im = hopper("BGR;24")
         assert im._repr_png_() is None
 
     def test_chunk_order(self, tmp_path):

Tests/test_image.py

@@ -105,6 +105,86 @@ def text(self, text):
         im._repr_pretty_(p, None)
         assert p.pretty_output == "<PIL.Image.Image image mode=L size=100x100>"
 
+    def test_repr_mimebundle(self):
+        im = Image.new("L", (100, 100))
+
+        # blank image should be most efficiently encoded as PNG
+        bundle = im._repr_mimebundle_()
+        with Image.open(io.BytesIO(bundle["image/png"])) as im2:
+            assert im2.format == "PNG"
+            assert_image_equal(im, im2)
+
+        # include pointless restriction
+        bundle = im._repr_mimebundle_(exclude=["test/plain"])
+        with Image.open(io.BytesIO(bundle["image/png"])) as im2:
+            assert im2.format == "PNG"
+            assert_image_equal(im, im2)
+
+        bundle = im._repr_mimebundle_(include=["image/png"])
+        with Image.open(io.BytesIO(bundle["image/png"])) as im2:
+            assert im2.format == "PNG"
+            assert_image_equal(im, im2)
+
+        # force jpeg to be selected
+        bundle = im._repr_mimebundle_(include=["image/jpeg"])
+        with Image.open(io.BytesIO(bundle["image/jpeg"])) as im2:
+            assert im2.format == "JPEG"
+            assert_image_equal(im, im2, 17)
+
+        # force jpeg to be selected in a different way
+        bundle = im._repr_mimebundle_(exclude=["image/png"])
+        with Image.open(io.BytesIO(bundle["image/jpeg"])) as im2:
+            assert im2.format == "JPEG"
+            assert_image_equal(im, im2, 17)
+
+        # make sure higher bit depths get converted down to 8BPC with warnings
+        high = Image.new("I;16", (100, 100), 65535)
+        bundle = high._repr_mimebundle_()
+        with Image.open(io.BytesIO(bundle["image/png"])) as im2:
+            assert im2.format == "PNG"
+            assert_image_equal(high.convert("I"), im2)
+
+        high = Image.new("F", (100, 100))
+        bundle = high._repr_mimebundle_()
+        assert bundle["image/jpeg"] is None and bundle["image/png"] is None
+
+        high = Image.new("I", (100, 100))
+        bundle = high._repr_mimebundle_()
+        with Image.open(io.BytesIO(bundle["image/png"])) as im2:
+            assert im2.format == "PNG"
+            assert_image_equal(high, im2)
+
+    def test_repr_mimebundle_hooks(self):
+        previous = Image._to_ipython_image
+        try:
+            Image.use_display_hook_features("auto")
+
+            # fmake sure large image gets scaled down with a warning
+            im = Image.new("L", [3000, 3000])
+            with pytest.warns(UserWarning):
+                bundle = im._repr_mimebundle_()
+
+            with Image.open(io.BytesIO(bundle["image/png"])) as im2:
+                assert im2.size == (1500, 1500)
+                assert_image_equal(im.resize(im2.size), im2)
+
+            # make sure common modes get converted without a warning
+            im = Image.new("LAB", (100, 100))
+            with pytest.warns(None) as record:
+                bundle = im._repr_mimebundle_()
+            assert len(record) == 0
+            with Image.open(io.BytesIO(bundle["image/png"])) as im2:
+                assert_image_equal(im.convert("RGB"), im2)
+
+            im = Image.new("HSV", (100, 100))
+            with pytest.warns(None) as record:
+                bundle = im._repr_mimebundle_()
+            assert len(record) == 0
+            with Image.open(io.BytesIO(bundle["image/png"])) as im2:
+                assert_image_equal(im.convert("RGB"), im2)
+        finally:
+            Image._to_ipython_image = previous
+
     def test_open_formats(self):
         PNGFILE = "Tests/images/hopper.png"
         JPGFILE = "Tests/images/hopper.jpg"

src/PIL/Image.py

@@ -72,6 +72,9 @@
 # Limit to around a quarter gigabyte for a 24-bit (3 bpp) image
 MAX_IMAGE_PIXELS = int(1024 * 1024 * 1024 // 4 // 3)
 
+# resize images much bigger than this when returning to IPython
+IPYTHON_RESIZE_THRESHOLD = 1200
+
 
 try:
     # If the _imaging C module is not present, Pillow will not load.
@@ -458,6 +461,185 @@
     return (a.scale, a.offset) if isinstance(a, _E) else (0, a)
 
 
+_IPYTHON_MODE_MAP = {
+    "La": "LA",
+    "LAB": "RGB",
+    "HSV": "RGB",
+    "RGBX": "RGB",
+    "RGBa": "RGBA",
+}
+
+_VALID_MODES_FOR_FORMAT = {
+    "JPEG": {"L", "RGB", "YCbCr", "CMYK"},
+    "PNG": {"1", "P", "L", "RGB", "RGBA", "LA", "PA", "I", "I;16"},
+}
+
+
+def _to_ipython_image(image):
+    """Transform image to something suitable for display in IPython/Jupyter.
+
+    See use_display_hook_features.
+    """
+    return image
+
+
+def use_display_hook_features(
+    default=None,
+    *,
+    I_mode=None,
+    F_mode=None,
+    I16_mode=None,
+    resize_threshold=None,
+    mode_map=None,
+):
+    """Set IPython/Jupyter display hook up based on parameters.
+
+    :param default: setting to use when unspecified
+    :param I_mode: what to do with I mode images
+    :param F_mode: what to do with F mode images
+    :param I16_mode: what to do with F mode images
+    :param resize_threshold: images with width or height much
+       larger than this will be resized
+    :param mode_map: a dictionary like _IPYTHON_MODE_MAP
+
+    when modes are set luminance values will be linearly rescaled such that;
+      'extrema': darkest value to black, brightest to white
+      'scale max': 0 value to black, brightest to white
+      'unit': 0 to black, 1 to white
+      'assume i16': 0 to black, 2^16-1 to white
+      'assume i32': 0 to black, 2^32-1 to white
+
+    All parameters can be specified as 'auto', which will cause the following
+    settings to be used:
+      I_mode, F_mode = 'scale_max'
+      I16_mode = 'assume i16'
+      resize_threshold = IPYTHON_RESIZE_THRESHOLD
+      mode_map = _IPYTHON_MODE_MAP
+    """
+
+    def identity(image):
+        "leave the image as is"
+        return image
+
+    def extrema_to_L(image):
+        # linearly transform extrema to fit in [0, 255]
+        # this should have a similar result as Image.histogram
+        lo, hi = image.getextrema()
+        warnings.warn(
+            f"converting image to 8BPC using black={lo}, white={hi} for display"
+        )
+        scale = 256 / (hi - lo) if lo != hi else 1
+        return image.point(lambda e: (e - lo) * scale).convert("L")
+
+    def scale_max_to_L(image):
+        _, hi = image.getextrema()
+        warnings.warn(f"converting image to 8BPC using black=0, white={hi} for display")
+        if hi > 0:
+            scale = 256 / hi
+            return image.point(lambda e: e * scale).convert("L")
+        return image
+
+    def unit_to_L(image):
+        warnings.warn("converting 8BPC for display, with range [0, 1]")
+        return image.point(lambda e: e * 255).convert("L")
+
+    def I16_to_L(image):
+        warnings.warn("converting 16BPC image to 8BPC for display")
+        # linearly transform max down to 255
+        return image.point(lambda e: e / 256).convert("L")
+
+    def I32_to_L(image):
+        warnings.warn("converting 32BPC image to 8BPC for display")
+        return image.point(lambda e: e / 2**24).convert("L")
+
+    def hook_for_mode(value, auto):
+        if value is None:
+            return identity
+        if value in "auto":
+            return auto
+        if value == "extrema":
+            return extrema_to_L
+        if value == "scale max":
+            return scale_max_to_L
+        if value == "unit":
+            return unit_to_L
+        if value == "assume i16":
+            return I16_to_L
+        if value == "assume i32":
+            return I32_to_L
+        assert isinstance(value, Callable)
+        return value
+
+    def threshold_resize(image):
+        # smaller images are quicker to process
+        factor = max(image.size) // resize_threshold
+        if factor > 1:
+            warnings.warn(
+                "scaling large image down to improve performance in IPython/Jupyter"
+            )
+            return image.reduce(factor)
+        return image
+
+    # decode arguments
+    handle_I = hook_for_mode(I_mode or default, scale_max_to_L)
+    handle_F = hook_for_mode(F_mode or default, scale_max_to_L)
+    handle_I16 = hook_for_mode(I16_mode or default, I16_to_L)
+
+    resize_threshold = resize_threshold or default
+    if resize_threshold is None:
+        resize_image = identity
+    else:
+        if resize_threshold == "auto":
+            resize_threshold = IPYTHON_RESIZE_THRESHOLD
+        assert isinstance(resize_threshold, int)
+        resize_image = threshold_resize
+
+    mode_map = mode_map or default
+    if mode_map == "auto":
+        mode_map = _IPYTHON_MODE_MAP
+    elif mode_map is None:
+        mode_map = {}
+    else:
+        assert isinstance(mode_map, dict)
+
+    # define our hook
+    def fn(image):
+        if image.mode == "I":
+            image = handle_I(image)
+        elif image.mode == "F":
+            image = handle_F(image)
+        elif image.mode == "I;16":
+            image = handle_I16(image)
+
+        image = resize_image(image)
+
+        # process remaining modes into things supported by writers
+        if image.mode in mode_map:
+            image = image.convert(mode_map[image.mode])
+
+        return image
+
+    # install it
+    global _to_ipython_image
+    _to_ipython_image = fn
+
+
+def _encode_ipython_image(image, image_format):
+    """Encode specfied image into something IPython/Jupyter supports.
+
+    :returns: bytes when valid, None when this is not valid
+    """
+    if image.mode not in _VALID_MODES_FOR_FORMAT.get(image_format, ()):
+        return None
+    b = io.BytesIO()
+    try:
+        image.save(b, image_format)
+    except Exception as e:
+        warnings.warn(f"failed to encode image as {image_format}: {e}")
+        return None
+    return b.getvalue()
+
+
 # --------------------------------------------------------------------
 # Implementation wrapper
 
@@ -632,32 +814,49 @@
             )
         )
 
-    def _repr_image(self, image_format, **kwargs):
-        """Helper function for iPython display hook.
+    def _repr_mimebundle_(self, include=None, exclude=None, **kwargs):
+        """iPython display hook that returns JPEG or PNG image as appropriate.
 
-        :param image_format: Image format.
-        :returns: image as bytes, saved into the given format.
+        :returns: iPython mimebundle
         """
-        b = io.BytesIO()
-        try:
-            self.save(b, image_format, **kwargs)
-        except Exception:
-            return None
-        return b.getvalue()
+        image = _to_ipython_image(self)
+
+        def encode(mimetype, image_format):
+            if include is not None and mimetype not in include:
+                return None
+            if exclude is not None and mimetype in exclude:
+                return None
+            return _encode_ipython_image(image, image_format)
+
+        jpeg = encode("image/jpeg", "JPEG")
+        png = encode("image/png", "PNG")
+
+        # prefer lossless format if it's not significantly larger
+        if jpeg and png:
+            # 1.125 and 2**18 used as they have nice binary representations
+            if len(png) < len(jpeg) * 1.125 + 2**18:
+                jpeg = None
+            else:
+                png = None
+
+        return {
+            "image/jpeg": jpeg,
+            "image/png": png,
+        }
 
     def _repr_png_(self):
         """iPython display hook support for PNG format.
 
         :returns: PNG version of the image as bytes
         """
-        return self._repr_image("PNG", compress_level=1)
+        return _encode_ipython_image(_to_ipython_image(self), "PNG")
 
     def _repr_jpeg_(self):
         """iPython display hook support for JPEG format.
 
         :returns: JPEG version of the image as bytes
         """
-        return self._repr_image("JPEG")
+        return _encode_ipython_image(_to_ipython_image(self), "JPEG")
 
     @property
     def __array_interface__(self):