Skip to content

Commit

Permalink
BUG: quick fix for pandas-dev#10989
Browse files Browse the repository at this point in the history
TST: add test case from Issue pandas-dev#10989
  • Loading branch information
jakevdp authored and jreback committed Oct 19, 2015
1 parent 5d953e3 commit d998337
Show file tree
Hide file tree
Showing 2 changed files with 28 additions and 0 deletions.
14 changes: 14 additions & 0 deletions pandas/tools/pivot.py
Original file line number Diff line number Diff line change
Expand Up @@ -159,6 +159,20 @@ def _add_margins(table, data, values, rows, cols, aggfunc):

grand_margin = _compute_grand_margin(data, values, aggfunc)

# categorical index or columns will fail below when 'All' is added
# here we'll convert all categorical indices to object
def convert_categorical(ind):
_convert = lambda ind: (ind.astype('object')
if ind.dtype.name == 'category' else ind)
if isinstance(ind, MultiIndex):
return ind.set_levels([_convert(lev) for lev in ind.levels])
else:
return _convert(ind)

table.index = convert_categorical(table.index)
if hasattr(table, 'columns'):
table.columns = convert_categorical(table.columns)

if not values and isinstance(table, Series):
# If there are no values and the table is a series, then there is only
# one column in the data. Compute grand margin and return it.
Expand Down
14 changes: 14 additions & 0 deletions pandas/tools/tests/test_pivot.py
Original file line number Diff line number Diff line change
Expand Up @@ -719,6 +719,20 @@ def test_crosstab_dropna(self):
('two', 'dull'), ('two', 'shiny')])
assert_equal(res.columns.values, m.values)

def test_categorical_margins(self):
# GH 10989
data = pd.DataFrame({'x': np.arange(8),
'y': np.arange(8) // 4,
'z': np.arange(8) % 2})
data.y = data.y.astype('category')
data.z = data.z.astype('category')
table = data.pivot_table('x', 'y', 'z', margins=True)
assert_equal(table.values, [[1, 2, 1.5],
[5, 6, 5.5],
[3, 4, 3.5]])



if __name__ == '__main__':
import nose
nose.runmodule(argv=[__file__, '-vvs', '-x', '--pdb', '--pdb-failure'],
Expand Down

0 comments on commit d998337

Please sign in to comment.