Skip to content

Commit 7a985f1

Browse files
committed
BUG: Better handling of invalid na_option argument for groupby.rank(pandas-dev#19499)
1 parent 415a01e commit 7a985f1

File tree

2 files changed

+20
-7
lines changed

2 files changed

+20
-7
lines changed

pandas/core/groupby/groupby.py

+3
Original file line numberDiff line numberDiff line change
@@ -1705,6 +1705,9 @@ def rank(self, method='average', ascending=True, na_option='keep',
17051705
-----
17061706
DataFrame with ranking of values within each group
17071707
"""
1708+
if na_option not in {'keep', 'top', 'bottom'}:
1709+
msg = "na_option must be one of 'keep', 'top', or 'bottom'"
1710+
raise ValueError(msg)
17081711
return self._cython_transform('rank', numeric_only=False,
17091712
ties_method=method, ascending=ascending,
17101713
na_option=na_option, pct=pct, axis=axis)

pandas/tests/groupby/test_rank.py

+17-7
Original file line numberDiff line numberDiff line change
@@ -252,14 +252,24 @@ def test_rank_object_raises(ties_method, ascending, na_option,
252252
with tm.assert_raises_regex(TypeError, "not callable"):
253253
df.groupby('key').rank(method=ties_method,
254254
ascending=ascending,
255-
na_option='bad', pct=pct)
255+
na_option=na_option, pct=pct)
256256

257-
with tm.assert_raises_regex(TypeError, "not callable"):
258-
df.groupby('key').rank(method=ties_method,
259-
ascending=ascending,
260-
na_option=True, pct=pct)
261257

262-
with tm.assert_raises_regex(TypeError, "not callable"):
258+
@pytest.mark.parametrize("na_option", [True, "bad", 1])
259+
@pytest.mark.parametrize("ties_method", [
260+
'average', 'min', 'max', 'first', 'dense'])
261+
@pytest.mark.parametrize("ascending", [True, False])
262+
@pytest.mark.parametrize("pct", [True, False])
263+
@pytest.mark.parametrize("vals", [
264+
['bar', 'bar', 'foo', 'bar', 'baz'],
265+
['bar', np.nan, 'foo', np.nan, 'baz'],
266+
[1, np.nan, 2, np.nan, 3]
267+
])
268+
def test_rank_naoption_raises(ties_method, ascending, na_option, pct, vals):
269+
df = DataFrame({'key': ['foo'] * 5, 'val': vals})
270+
msg = "na_option must be one of 'keep', 'top', or 'bottom'"
271+
272+
with tm.assert_raises_regex(TypeError, msg):
263273
df.groupby('key').rank(method=ties_method,
264274
ascending=ascending,
265-
na_option=na_option, pct=pct)
275+
na_option='bad', pct=pct)

0 commit comments

Comments
 (0)