ScalarMappable.to_rgba can return uint8 instead of float64

svn path=/trunk/matplotlib/; revision=4328

lib/matplotlib/cm.py

@@ -45,14 +45,18 @@ def set_colorbar(self, im, ax):
         'set the colorbar image and axes associated with mappable'
         self.colorbar = im, ax
 
-    def to_rgba(self, x, alpha=1.0):
+    def to_rgba(self, x, alpha=1.0, bytes=False):
         '''Return a normalized rgba array corresponding to x.
         If x is already an rgb or rgba array, return it unchanged.
         '''
-        if hasattr(x, 'shape') and len(x.shape)>2: return x
+        try:
+            if x.ndim == 3 and (x.shape[2] == 3 or x.shape[2] == 4):
+                return x
+        except AttributeError:
+            pass
         x = ma.asarray(x)
         x = self.norm(x)
-        x = self.cmap(x, alpha)
+        x = self.cmap(x, alpha=alpha, bytes=bytes)
         return x
 
     def set_array(self, A):

lib/matplotlib/colors.py

@@ -406,7 +406,7 @@ def __init__(self, name, N=256):
         self._isinit = False
 
 
-    def __call__(self, X, alpha=1.0):
+    def __call__(self, X, alpha=1.0, bytes=False):
         """
         X is either a scalar or an array (of any dimension).
         If scalar, a tuple of rgba values is returned, otherwise
@@ -415,6 +415,8 @@ def __call__(self, X, alpha=1.0):
         If they are floating point, then they must be in the
         interval (0.0, 1.0).
         Alpha must be a scalar.
+        If bytes is False, the rgba values will be floats on a
+        0-1 scale; if True, they will be uint8, 0-255.
         """
 
         if not self._isinit: self._init()
@@ -439,7 +441,11 @@ def __call__(self, X, alpha=1.0):
         npy.putmask(xa, xa<0, self._i_under)
         if mask_bad is not None and mask_bad.shape == xa.shape:
             npy.putmask(xa, mask_bad, self._i_bad)
-        rgba = self._lut[xa]
+        if bytes:
+            lut = (self._lut * 255).astype(npy.uint8)
+        else:
+            lut = self._lut
+        rgba = lut[xa]
         if vtype == 'scalar':
             rgba = tuple(rgba[0,:])
         return rgba