From 7ea45833f457e4320b3f0ada02d484b9e2d207d3 Mon Sep 17 00:00:00 2001 From: sinhrks Date: Sat, 24 May 2014 16:22:01 +0900 Subject: [PATCH] BUG: boxplot returns incorrect dict --- pandas/tests/test_graphics.py | 45 ++++++++++++++++++++++++++++++++++- pandas/tools/plotting.py | 14 +++++++---- 2 files changed, 53 insertions(+), 6 deletions(-) diff --git a/pandas/tests/test_graphics.py b/pandas/tests/test_graphics.py index 39067096cfd25..de726e670d958 100644 --- a/pandas/tests/test_graphics.py +++ b/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 diff --git a/pandas/tools/plotting.py b/pandas/tools/plotting.py index 4d5d8d773a469..00f7a61870369 100644 --- a/pandas/tools/plotting.py +++ b/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]