From 1583a04a0c13d8d1f8f17874068dfa23f89310cb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jonas=20H=C3=B6rsch?= Date: Mon, 22 May 2023 19:28:14 +0200 Subject: [PATCH] Add callable filters to isin selector (#16) --- CHANGELOG.rst | 2 ++ src/pandas_indexing/selectors.py | 9 +++++++-- tests/test_selectors.py | 3 +++ 3 files changed, 12 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 46aa317..e8b4905 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -3,6 +3,8 @@ Changelog ========= +* :func:`~selectors.isin` accepts callable filters :pull:`16`, f.ex. + ``df.loc[isin(year=lambda s: s>2000)]`` * New function :func:`~core.concat` makes concatenation level aware :pull:`14` v0.2.5 (2023-05-04) diff --git a/src/pandas_indexing/selectors.py b/src/pandas_indexing/selectors.py index b167a48..2223db2 100644 --- a/src/pandas_indexing/selectors.py +++ b/src/pandas_indexing/selectors.py @@ -84,8 +84,12 @@ def __call__(self, df): if self.ignore_missing_levels: filters = {k: v for k, v in filters.items() if k in index.names} - tests = (index.isin(np.atleast_1d(v), level=k) for k, v in filters.items()) - return reduce(and_, tests) + def apply_filter(value, level): + if callable(value): + return value(index.get_level_values(level)) + return index.isin(np.atleast_1d(value), level=level) + + return reduce(and_, (apply_filter(v, k) for k, v in filters.items())) def isin( @@ -101,6 +105,7 @@ def isin( If set, levels missing in data index will be ignored **filters Filter to apply on given levels (lists are ``or`` ed, levels are ``and`` ed) + Callables are evaluated on the index level values. Returns ------- diff --git a/tests/test_selectors.py b/tests/test_selectors.py index 550d186..70f3f23 100644 --- a/tests/test_selectors.py +++ b/tests/test_selectors.py @@ -17,6 +17,9 @@ def test_isin_mseries(mseries: Series): sel = isin(num=[2, 3], str="foo") assert_series_equal(mseries.loc[sel], mseries.iloc[[1]]) + sel = isin(num=lambda s: s > 1) + assert_series_equal(mseries.loc[sel], mseries.iloc[[1, 2]]) + def test_isin_ignore_missing_levels(mseries: Series): sel = isin(str="foo", bla=False)