Allow imsave() to handle RGB(A) input arrays

WeatherGod · WeatherGod · commit 70358acc9e08 · 2012-09-09T15:45:11.000-04:00
Thanks to Jostein Be Floystad for the original patch!
diff --git a/lib/matplotlib/image.py b/lib/matplotlib/image.py
@@ -1243,7 +1243,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
@@ -1266,7 +1266,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)
