| 1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768 |
- import numpy as np
- import pytest
- import pandas as pd
- from pandas.util.testing import assert_frame_equal, assert_series_equal
- @pytest.fixture(params=[['inner'], ['inner', 'outer']])
- def frame(request):
- levels = request.param
- df = pd.DataFrame({'outer': ['a', 'a', 'a', 'b', 'b', 'b'],
- 'inner': [1, 2, 3, 1, 2, 3],
- 'A': np.arange(6),
- 'B': ['one', 'one', 'two', 'two', 'one', 'one']})
- if levels:
- df = df.set_index(levels)
- return df
- @pytest.fixture()
- def series():
- df = pd.DataFrame({'outer': ['a', 'a', 'a', 'b', 'b', 'b'],
- 'inner': [1, 2, 3, 1, 2, 3],
- 'A': np.arange(6),
- 'B': ['one', 'one', 'two', 'two', 'one', 'one']})
- s = df.set_index(['outer', 'inner', 'B'])['A']
- return s
- @pytest.mark.parametrize('key_strs,groupers', [
- ('inner', # Index name
- pd.Grouper(level='inner')
- ),
- (['inner'], # List of index name
- [pd.Grouper(level='inner')]
- ),
- (['B', 'inner'], # Column and index
- ['B', pd.Grouper(level='inner')]
- ),
- (['inner', 'B'], # Index and column
- [pd.Grouper(level='inner'), 'B'])])
- def test_grouper_index_level_as_string(frame, key_strs, groupers):
- result = frame.groupby(key_strs).mean()
- expected = frame.groupby(groupers).mean()
- assert_frame_equal(result, expected)
- @pytest.mark.parametrize('levels', [
- 'inner', 'outer', 'B',
- ['inner'], ['outer'], ['B'],
- ['inner', 'outer'], ['outer', 'inner'],
- ['inner', 'outer', 'B'], ['B', 'outer', 'inner']
- ])
- def test_grouper_index_level_as_string_series(series, levels):
- # Compute expected result
- if isinstance(levels, list):
- groupers = [pd.Grouper(level=lv) for lv in levels]
- else:
- groupers = pd.Grouper(level=levels)
- expected = series.groupby(groupers).mean()
- # Compute and check result
- result = series.groupby(levels).mean()
- assert_series_equal(result, expected)
|
