Add ability to render bar charts and legends for bar charts (#135)

* Add ability to render bar chart legends

Bar charts in matplotlib consist of a set of rectangles. This commit
finds the first rectange of each label, and provides a \addlegendimage
with the appropriate draw options.

* Removed draw_rectanges, added barchart test with no legend

Author: dougnd

matplotlib2tikz/patch.py

@@ -51,14 +51,36 @@ def draw_patchcollection(data, obj):
 def _draw_rectangle(data, obj, draw_options):
     '''Return the PGFPlots code for rectangles.
     '''
-    if not data['draw rectangles']:
+
+    # Objects with labels are plot objects (from bar charts, etc).
+    # Even those without labels explicitly set have a label of
+    # "_nolegend_".  Everything else should be skipped because
+    # they likely correspong to axis/legend objects which are
+    # handled by PGFPlots
+    label = obj.get_label()
+    if label == '':
         return data, []
 
+    # get real label, bar charts by default only give rectangles
+    # labels of "_nolegend_"
+    # See http://stackoverflow.com/questions/35881290/how-to-get-the-label-on-bar-plot-stacked-bar-plot-in-matplotlib
+    handles,labels = obj.axes.get_legend_handles_labels()
+    labelsFound = [label for h,label in zip(handles, labels) if obj in h.get_children()]
+    if len(labelsFound) == 1:
+        label = labelsFound[0]
+
+    legend = ''
+    if label != '_nolegend_' and label not in data['rectangle_legends']:
+        data['rectangle_legends'].add(label)
+        legend = ('\\addlegendimage{ybar,ybar legend,%s};\n'
+                 ) % (','.join(draw_options))
+
     left_lower_x = obj.get_x()
     left_lower_y = obj.get_y()
-    cont = ('\draw[%s] (axis cs:%.15g,%.15g) '
+    cont = ('%s\draw[%s] (axis cs:%.15g,%.15g) '
             'rectangle (axis cs:%.15g,%.15g);\n'
-            ) % (','.join(draw_options),
+            ) % (legend,
+                 ','.join(draw_options),
                  left_lower_x,
                  left_lower_y,
                  left_lower_x + obj.get_width(),

matplotlib2tikz/save.py

@@ -22,7 +22,6 @@ def save(filepath,
          textsize=10.0,
          tex_relative_path_to_data=None,
          strict=False,
-         draw_rectangles=False,
          wrap=True,
          extra=None,
          dpi=None,
@@ -75,16 +74,6 @@ def save(filepath,
                    can decide where to put the ticks.
     :type strict: bool
 
-    :param draw_rectangles: Whether or not to draw Rectangle objects.
-                            You normally don't want that as legend, axes, and
-                            other entities which are natively taken care of by
-                            PGFPlots are represented as rectangles in
-                            matplotlib. Some plot types (such as bar plots)
-                            cannot otherwise be represented though.
-                            Don't expect working or clean output when using
-                            this option.
-    :type draw_rectangles: bool
-
     :param wrap: Whether ``'\\begin{tikzpicture}'`` and
                  ``'\\end{tikzpicture}'`` will be written. One might need to
                  provide custom arguments to the environment (eg. scale= etc.).
@@ -121,11 +110,14 @@ def save(filepath,
     data['output dir'] = os.path.dirname(filepath)
     data['base name'] = os.path.splitext(os.path.basename(filepath))[0]
     data['strict'] = strict
-    data['draw rectangles'] = draw_rectangles
     data['tikz libs'] = set()
     data['pgfplots libs'] = set()
     data['font size'] = textsize
     data['custom colors'] = {}
+    # rectangle_legends is used to keep track of which rectangles have already
+    # had \addlegendimage added. There should be only one \addlegenimage per 
+    # bar chart data series.
+    data['rectangle_legends'] = set()
     if extra:
         data['extra axis options'] = extra.copy()
     else:

test/test_hashes.py

@@ -32,7 +32,6 @@ def test_hash(name):
         figureheight='7.5cm',
         show_info=True,
         strict=True,
-        draw_rectangles=True
         )
 
     # save reference figure

test/testfunctions/__init__.py

@@ -4,6 +4,8 @@
   'annotate',
   'basic_sin',
   'boxplot',
+  'barchart_legend',
+  'barchart',
   'dual_axis',
   'errorband',
   'errorbar',

test/testfunctions/barchart.py

@@ -0,0 +1,31 @@
+# -*- coding: utf-8 -*-
+""" Bar Chart test
+
+This tests plots a simple bar chart.  Bar charts are plotted as
+rectangle patches witch are difficult to tell from other rectangle
+patches that should not be plotted in PGFPlots (e.g. axis, legend)
+
+"""
+desc = 'Bar Chart'
+phash = '5f09a9e6b3728742'
+
+def plot():
+    import matplotlib.pyplot as plt
+    import numpy as np
+
+    # plot data
+    fig = plt.figure()
+    ax = fig.add_subplot(111)
+
+    x = np.arange(3)
+    y1 = [1, 2, 3]
+    y2 = [3, 2, 4]
+    y3 = [5, 3, 1]
+    w = 0.25
+
+    ax.bar(x-w, y1, w, color='b', align='center')
+    ax.bar(x,   y2, w, color='g', align='center')
+    ax.bar(x+w, y3, w, color='r', align='center')
+
+    return fig
+

test/testfunctions/barchart_legend.py

@@ -0,0 +1,36 @@
+# -*- coding: utf-8 -*-
+""" Bar Chart Legend test
+
+This tests plots a simple bar chart.  Bar charts are plotted as
+rectangle patches witch are difficult to tell from other rectangle
+patches that should not be plotted in PGFPlots (e.g. axis, legend)
+
+This also tests legends on barcharts. Which are difficult because
+in PGFPlots, they have no \\addplot, and thus legend must be 
+manually added.
+
+"""
+desc = 'Bar Chart'
+phash = '5f09a9e633728dc4'
+
+def plot():
+    import matplotlib.pyplot as plt
+    import numpy as np
+
+    # plot data
+    fig = plt.figure()
+    ax = fig.add_subplot(111)
+
+    x = np.arange(3)
+    y1 = [1, 2, 3]
+    y2 = [3, 2, 4]
+    y3 = [5, 3, 1]
+    w = 0.25
+
+    ax.bar(x-w, y1, w, color='b', align='center', label='Data 1')
+    ax.bar(x,   y2, w, color='g', align='center', label='Data 2')
+    ax.bar(x+w, y3, w, color='r', align='center', label='Data 3')
+    ax.legend()
+
+    return fig
+