Merge pull request matplotlib#2389 from bjcohen/patch-1

tacaswell · tacaswell · commit be3da1aed59c · 2014-01-31T20:07:19.000-08:00
table.py: fix issue when specifying both column header text and color
diff --git a/lib/matplotlib/table.py b/lib/matplotlib/table.py
@@ -495,26 +495,29 @@ def table(ax,
     if colWidths is None:
         colWidths = [1.0 / cols] * cols
 
-    # Check row and column labels
+    # Fill in missing information for column
+    # and row labels
     rowLabelWidth = 0
     if rowLabels is None:
         if rowColours is not None:
-            rowLabels = [''] * cols
+            rowLabels = [''] * rows
             rowLabelWidth = colWidths[0]
     elif rowColours is None:
         rowColours = 'w' * rows
 
     if rowLabels is not None:
         assert len(rowLabels) == rows
 
-    offset = 0
+    # If we have column labels, need to shift
+    # the text and colour arrays down 1 row
+    offset = 1
     if colLabels is None:
         if colColours is not None:
-            colLabels = [''] * rows
-            offset = 1
+            colLabels = [''] * cols
+        else:
+            offset = 0
     elif colColours is None:
         colColours = 'w' * cols
-        offset = 1
 
     if rowLabels is not None:
         assert len(rowLabels) == rows
diff --git a/lib/matplotlib/tests/baseline_images/test_table/table_labels.png b/lib/matplotlib/tests/baseline_images/test_table/table_labels.png
diff --git a/lib/matplotlib/tests/test_table.py b/lib/matplotlib/tests/test_table.py
@@ -41,3 +41,41 @@ def test_zorder():
               zorder=4,
               )
     plt.yticks([])
+
+
+@image_comparison(baseline_images=['table_labels'],
+                  extensions=['png'])
+def test_label_colours():
+    dim = 3
+
+    c = np.linspace(0, 1, dim)
+    colours = plt.cm.RdYlGn(c)
+    cellText = [['1'] * dim] * dim
+
+    fig = plt.figure()
+
+    ax1 = fig.add_subplot(4, 1, 1)
+    ax1.axis('off')
+    ax1.table(cellText=cellText,
+              rowColours=colours,
+              loc='best')
+
+    ax2 = fig.add_subplot(4, 1, 2)
+    ax2.axis('off')
+    ax2.table(cellText=cellText,
+              rowColours=colours,
+              rowLabels=['Header'] * dim,
+              loc='best')
+
+    ax3 = fig.add_subplot(4, 1, 3)
+    ax3.axis('off')
+    ax3.table(cellText=cellText,
+              colColours=colours,
+              loc='best')
+
+    ax4 = fig.add_subplot(4, 1, 4)
+    ax4.axis('off')
+    ax4.table(cellText=cellText,
+              colColours=colours,
+              colLabels=['Header'] * dim,
+              loc='best')
