| 104 | @pytest.mark.parametrize("foo", (1, None)) |
| 105 | @pytest.mark.parametrize("bar", (2, None)) |
| 106 | def test_rolling_apply(df, pdf, window, raw, foo, bar): |
| 107 | def my_sum(vals, foo_=None, *, bar_=None): |
| 108 | assert foo_ == foo |
| 109 | assert bar_ == bar |
| 110 | if raw: |
| 111 | assert isinstance(vals, np.ndarray) |
| 112 | else: |
| 113 | assert isinstance(vals, pd.Series) |
| 114 | return vals.sum() |
| 115 | |
| 116 | kwargs = dict(raw=raw, args=(foo,), kwargs=dict(bar_=bar)) |
| 117 | |
| 118 | result = df.rolling(window).apply(my_sum, **kwargs) |
| 119 | expected = pdf.rolling(window).apply(my_sum, **kwargs) |
| 120 | assert_eq(result, expected) |
| 121 | |
| 122 | result = df.rolling(window).apply(my_sum, **kwargs)["foo"] |
| 123 | expected = pdf.rolling(window).apply(my_sum, **kwargs)["foo"] |
| 124 | assert_eq(result, expected) |
| 125 | |
| 126 | # simplify up disabled for `apply`, function may access other columns |
| 127 | q = df.rolling(window).apply(my_sum, **kwargs)["foo"].simplify() |
| 128 | eq = df["foo"].rolling(window).apply(my_sum, **kwargs).simplify() |
| 129 | assert q._name == eq._name |
| 130 | |
| 131 | |
| 132 | def test_rolling_one_element_window(df, pdf): |