Created _repr_png_ and _repr_html_ for Colormap objects.

bdice · bdice · commit 1cfb18bbe864 · 2020-07-11T12:22:45.000-05:00
diff --git a/lib/matplotlib/colors.py b/lib/matplotlib/colors.py
@@ -65,11 +65,14 @@
 .. _xkcd color survey: https://xkcd.com/color/rgb/
 """
 
+import base64
 from collections.abc import Sized
 import functools
+import io
 import itertools
 from numbers import Number
 import re
+from PIL import Image
 
 import numpy as np
 import matplotlib.cbook as cbook
@@ -691,6 +694,25 @@ def reversed(self, name=None):
         """
         raise NotImplementedError()
 
+    def _repr_png_(self):
+        """Generate a PNG representation of the Colormap."""
+        IMAGE_SIZE = (400, 50)
+        X = np.tile(np.linspace(0, 1, IMAGE_SIZE[0]), (IMAGE_SIZE[1], 1))
+        pixels = self(X, bytes=True)
+        png_bytes = io.BytesIO()
+        Image.fromarray(pixels).save(png_bytes, format='png')
+        return png_bytes.getvalue()
+
+    def _repr_html_(self):
+        """Generate an HTML representation of the Colormap."""
+        png_bytes = self._repr_png_()
+        png_base64 = base64.b64encode(png_bytes).decode('ascii')
+        return ('<strong>' + self.name + '</strong>' +
+                '<img ' +
+                'alt="' + self.name + ' color map" ' +
+                'title="' + self.name + '"' +
+                'src="data:image/png;base64,' + png_base64 + '">')
+
 
 class LinearSegmentedColormap(Colormap):
     """
