Use list comprehensions to ravel

patniharshit · patniharshit · commit 2573610cd817 · 2017-08-27T15:00:04.000+05:30
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -5450,11 +5450,12 @@ def pcolor(self, *args, **kwargs):
         verts = xy.reshape((npoly, 5, 2))
 
         if isinstance(norm, mcolors.BivariateNorm):
-            C0 = C[0]
-            C1 = C[1]
-            C0 = compress(ravelmask, ma.filled(C0[0:Ny - 1, 0:Nx - 1]).ravel())
-            C1 = compress(ravelmask, ma.filled(C1[0:Ny - 1, 0:Nx - 1]).ravel())
-            C = np.array([C0, C1])
+            C = np.array([
+                            compress(
+                                ravelmask,
+                                ma.filled(c[0:Ny - 1, 0:Nx - 1]).ravel()
+                            ) for c in C
+                        ])
         else:
             C = compress(ravelmask, ma.filled(C[0:Ny - 1, 0:Nx - 1]).ravel())
 
@@ -5629,7 +5630,7 @@ def pcolormesh(self, *args, **kwargs):
 
         # convert to one dimensional arrays if univariate
         if isinstance(norm, mcolors.BivariateNorm):
-            C = np.asarray([C[0].ravel(), C[1].ravel()])
+            C = np.asarray([c.ravel() for c in C])
         else:
             C = C.ravel()
 
