Merge pull request #7225 from sinhrks/boxbug

BUG: boxplot returns incorrect dict

Author: jreback

pandas/tests/test_graphics.py

@@ -10,7 +10,8 @@
 from datetime import datetime, date
 
 from pandas import Series, DataFrame, MultiIndex, PeriodIndex, date_range
-from pandas.compat import range, lrange, StringIO, lmap, lzip, u, zip
+from pandas.compat import (range, lrange, StringIO, lmap, lzip, u, zip,
+                           iteritems, OrderedDict)
 from pandas.util.decorators import cache_readonly
 import pandas.core.common as com
 import pandas.util.testing as tm
@@ -2245,6 +2246,48 @@ def test_grouped_hist(self):
         with tm.assertRaises(AttributeError):
             plotting.grouped_hist(df.A, by=df.C, foo='bar')
 
+    def _check_box_dict(self, returned, return_type,
+                        expected_klass, expected_keys):
+        self.assertTrue(isinstance(returned, OrderedDict))
+        self.assertEqual(sorted(returned.keys()), sorted(expected_keys))
+        for key, value in iteritems(returned):
+            self.assertTrue(isinstance(value, expected_klass))
+            # check returned dict has correct mapping
+            if return_type == 'axes':
+                self.assertEqual(value.get_title(), key)
+            elif return_type == 'both':
+                self.assertEqual(value.ax.get_title(), key)
+            elif return_type == 'dict':
+                line = value['medians'][0]
+                self.assertEqual(line.get_axes().get_title(), key)
+            else:
+                raise AssertionError
+
+    @slow
+    def test_grouped_box_return_type(self):
+        import matplotlib.axes
+
+        df = self.hist_df
+
+        columns2 = 'X B C D A G Y N Q O'.split()
+        df2 = DataFrame(random.randn(50, 10), columns=columns2)
+        categories2 = 'A B C D E F G H I J'.split()
+        df2['category'] = tm.choice(categories2, size=50)
+
+        types = {'dict': dict, 'axes': matplotlib.axes.Axes, 'both': tuple}
+        for t, klass in iteritems(types):
+            returned = df.groupby('classroom').boxplot(return_type=t)
+            self._check_box_dict(returned, t, klass, ['A', 'B', 'C'])
+
+            returned = df.boxplot(by='classroom', return_type=t)
+            self._check_box_dict(returned, t, klass, ['height', 'weight', 'category'])
+
+            returned = df2.groupby('category').boxplot(return_type=t)
+            self._check_box_dict(returned, t, klass, categories2)
+
+            returned = df2.boxplot(by='category', return_type=t)
+            self._check_box_dict(returned, t, klass, columns2)
+
     @slow
     def test_grouped_box_layout(self):
         df = self.hist_df

pandas/tools/plotting.py

@@ -2363,12 +2363,17 @@ def plot_group(grouped, ax):
         if return_type is None:
             ret = axes
         if return_type == 'axes':
-            ret = dict((k, ax) for k, ax in zip(d.keys(), axes))
+            ret = compat.OrderedDict()
+            axes = _flatten(axes)[:len(d)]
+            for k, ax in zip(d.keys(), axes):
+                ret[k] = ax
         elif return_type == 'dict':
             ret = d
         elif return_type == 'both':
-            ret = dict((k, BP(ax=ax, lines=line)) for
-                       (k, line), ax in zip(d.items(), axes))
+            ret = compat.OrderedDict()
+            axes = _flatten(axes)[:len(d)]
+            for (k, line), ax in zip(d.items(), axes):
+                ret[k] = BP(ax=ax, lines=line)
     else:
         if layout is not None:
             raise ValueError("The 'layout' keyword is not supported when "
@@ -2723,7 +2728,7 @@ def boxplot_frame_groupby(grouped, subplots=True, column=None, fontsize=None,
                             sharex=False, sharey=True)
         axes = _flatten(axes)
 
-        ret = {}
+        ret = compat.OrderedDict()
         for (key, group), ax in zip(grouped, axes):
             d = group.boxplot(ax=ax, column=column, fontsize=fontsize,
                               rot=rot, grid=grid, **kwds)
@@ -2804,7 +2809,6 @@ def _grouped_plot_by_column(plotf, data, columns=None, by=None,
     ravel_axes = _flatten(axes)
 
     out_dict = compat.OrderedDict()
-
     for i, col in enumerate(columns):
         ax = ravel_axes[i]
         gp_col = grouped[col]