Merge pull request #1224 from WeatherGod/imshow_alpha

mdboom · mdboom · commit 5b2a72a0679a · 2012-09-12T18:29:01.000-07:00
Imsave alpha
diff --git a/lib/matplotlib/image.py b/lib/matplotlib/image.py
@@ -1246,7 +1246,7 @@ def imsave(fname, arr, vmin=None, vmax=None, cmap=None, format=None,
         If *format* is *None* and *fname* is a string, the output
         format is deduced from the extension of the filename.
       *arr*:
-        A 2D array.
+        An MxN (luminance), MxNx3 (RGB) or MxNx4 (RGBA) array.
     Keyword arguments:
       *vmin*/*vmax*: [ None | scalar ]
         *vmin* and *vmax* set the color scaling for the image by fixing the
@@ -1269,7 +1269,7 @@ def imsave(fname, arr, vmin=None, vmax=None, cmap=None, format=None,
     from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
     from matplotlib.figure import Figure
 
-    figsize = [x / float(dpi) for x in arr.shape[::-1]]
+    figsize = [x / float(dpi) for x in (arr.shape[1], arr.shape[0])]
     fig = Figure(figsize=figsize, dpi=dpi, frameon=False)
     canvas = FigureCanvas(fig)
     im = fig.figimage(arr, cmap=cmap, vmin=vmin, vmax=vmax, origin=origin)
diff --git a/lib/matplotlib/tests/test_image.py b/lib/matplotlib/tests/test_image.py
@@ -128,6 +128,26 @@ def test_imsave():
 
     assert_array_equal(arr_dpi1, arr_dpi100)
 
+def test_imsave_color_alpha():
+    # The goal is to test that imsave will accept arrays with ndim=3 where
+    # the third dimension is color and alpha without raising any exceptions
+    from numpy import random
+    random.seed(1)
+    data = random.rand(256, 128, 4)
+
+    buff = io.BytesIO()
+    plt.imsave(buff, data)
+
+    buff.seek(0)
+    arr_buf = plt.imread(buff)
+
+    assert arr_buf.shape == data.shape
+
+    # Unfortunately, the AGG process "flattens" the RGBA data
+    # into an equivalent RGB data with no transparency. So we
+    # Can't directly compare the arrays like we could in some
+    # other imsave tests.
+
 @image_comparison(baseline_images=['image_clip'])
 def test_image_clip():
     from math import pi
