|
7 | 7 | from pandas.core.dtypes.common import is_number
|
8 | 8 |
|
9 | 9 | import pandas as pd
|
10 |
| -from pandas import DataFrame, Index, MultiIndex, Series, isna, timedelta_range |
| 10 | +from pandas import DataFrame, Index, MultiIndex, Series, concat, isna, timedelta_range |
11 | 11 | import pandas._testing as tm
|
12 | 12 | from pandas.core.base import SpecificationError
|
13 | 13 |
|
@@ -827,3 +827,75 @@ def test_apply_to_timedelta(self):
|
827 | 827 | b = Series(list_of_strings).apply(pd.to_timedelta) # noqa
|
828 | 828 | # Can't compare until apply on a Series gives the correct dtype
|
829 | 829 | # assert_series_equal(a, b)
|
| 830 | + |
| 831 | + |
| 832 | +@pytest.mark.parametrize( |
| 833 | + "ops, names", |
| 834 | + [ |
| 835 | + ([np.sum], ["sum"]), |
| 836 | + ([np.sum, np.mean], ["sum", "mean"]), |
| 837 | + (np.array([np.sum]), ["sum"]), |
| 838 | + (np.array([np.sum, np.mean]), ["sum", "mean"]), |
| 839 | + ], |
| 840 | +) |
| 841 | +@pytest.mark.parametrize("how", ["agg", "apply"]) |
| 842 | +def test_apply_listlike_reducer(string_series, ops, names, how): |
| 843 | + # GH 39140 |
| 844 | + expected = Series({name: op(string_series) for name, op in zip(names, ops)}) |
| 845 | + expected.name = "series" |
| 846 | + result = getattr(string_series, how)(ops) |
| 847 | + tm.assert_series_equal(result, expected) |
| 848 | + |
| 849 | + |
| 850 | +@pytest.mark.parametrize( |
| 851 | + "ops", |
| 852 | + [ |
| 853 | + {"A": np.sum}, |
| 854 | + {"A": np.sum, "B": np.mean}, |
| 855 | + Series({"A": np.sum}), |
| 856 | + Series({"A": np.sum, "B": np.mean}), |
| 857 | + ], |
| 858 | +) |
| 859 | +@pytest.mark.parametrize("how", ["agg", "apply"]) |
| 860 | +def test_apply_dictlike_reducer(string_series, ops, how): |
| 861 | + # GH 39140 |
| 862 | + expected = Series({name: op(string_series) for name, op in ops.items()}) |
| 863 | + expected.name = string_series.name |
| 864 | + result = getattr(string_series, how)(ops) |
| 865 | + tm.assert_series_equal(result, expected) |
| 866 | + |
| 867 | + |
| 868 | +@pytest.mark.parametrize( |
| 869 | + "ops, names", |
| 870 | + [ |
| 871 | + ([np.sqrt], ["sqrt"]), |
| 872 | + ([np.abs, np.sqrt], ["absolute", "sqrt"]), |
| 873 | + (np.array([np.sqrt]), ["sqrt"]), |
| 874 | + (np.array([np.abs, np.sqrt]), ["absolute", "sqrt"]), |
| 875 | + ], |
| 876 | +) |
| 877 | +def test_apply_listlike_transformer(string_series, ops, names): |
| 878 | + # GH 39140 |
| 879 | + with np.errstate(all="ignore"): |
| 880 | + expected = concat([op(string_series) for op in ops], axis=1) |
| 881 | + expected.columns = names |
| 882 | + result = string_series.apply(ops) |
| 883 | + tm.assert_frame_equal(result, expected) |
| 884 | + |
| 885 | + |
| 886 | +@pytest.mark.parametrize( |
| 887 | + "ops", |
| 888 | + [ |
| 889 | + {"A": np.sqrt}, |
| 890 | + {"A": np.sqrt, "B": np.exp}, |
| 891 | + Series({"A": np.sqrt}), |
| 892 | + Series({"A": np.sqrt, "B": np.exp}), |
| 893 | + ], |
| 894 | +) |
| 895 | +def test_apply_dictlike_transformer(string_series, ops): |
| 896 | + # GH 39140 |
| 897 | + with np.errstate(all="ignore"): |
| 898 | + expected = concat({name: op(string_series) for name, op in ops.items()}) |
| 899 | + expected.name = string_series.name |
| 900 | + result = string_series.apply(ops) |
| 901 | + tm.assert_series_equal(result, expected) |
0 commit comments