From be1972cc0912f971fa62af962d457ecd473d3b62 Mon Sep 17 00:00:00 2001 From: Giovanni Lanzani Date: Mon, 15 Jun 2020 08:54:41 +0200 Subject: [PATCH] TST Add test for rolling window, see GH 34605 --- pandas/tests/window/test_rolling.py | 33 +++++++++++++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/pandas/tests/window/test_rolling.py b/pandas/tests/window/test_rolling.py index f9b0e6856337b..8d72e2cb92ca9 100644 --- a/pandas/tests/window/test_rolling.py +++ b/pandas/tests/window/test_rolling.py @@ -663,3 +663,36 @@ def test_iter_rolling_datetime(expected, expected_index, window): for (expected, actual) in zip(expected, ser.rolling(window)): tm.assert_series_equal(actual, expected) + + +@pytest.mark.parametrize( + "grouping,_index", + [ + ( + {"level": 0}, + pd.MultiIndex.from_tuples( + [(0, 0), (0, 0), (1, 1), (1, 1), (1, 1)], names=[None, None] + ), + ), + ( + {"by": "X"}, + pd.MultiIndex.from_tuples( + [(0, 0), (1, 0), (2, 1), (3, 1), (4, 1)], names=["X", None] + ), + ), + ], +) +def test_rolling_positional_argument(grouping, _index, raw): + # GH 34605 + + def scaled_sum(*args): + if len(args) < 2: + raise ValueError("The function needs two arguments") + array, scale = args + return array.sum() / scale + + df = DataFrame(data={"X": range(5)}, index=[0, 0, 1, 1, 1]) + + expected = DataFrame(data={"X": [0.0, 0.5, 1.0, 1.5, 2.0]}, index=_index) + result = df.groupby(**grouping).rolling(1).apply(scaled_sum, raw=raw, args=(2,)) + tm.assert_frame_equal(result, expected)