diff --git a/docs/reference.md b/docs/reference.md index 80b2144..bed6fb9 100644 --- a/docs/reference.md +++ b/docs/reference.md @@ -435,6 +435,21 @@ False ``` +## YieldIf +Takes a function to map values (optional, by default there's no mapping) and a key. If key is false, value will not be yielded. Key is optional, default is `bool` + +```py +>>> range(5) | YieldIf(lambda x: x * 100) | Pipe(list) +[100, 200, 300, 400] +>>> range(5) | YieldIf(lambda x: x * 100, key=lambda x: x % 2 == 0) | Pipe(list) +[0, 200, 400] +>>> range(5) | YieldIf(key=lambda x: x % 2 == 0) | Pipe(list) +[0, 2, 4] +>>> range(5) | YieldIf() | Pipe(list) +[1, 2, 3, 4] + +``` + ## GetItem ```py diff --git a/pipe21.py b/pipe21.py index d9a9291..d2bd94e 100644 --- a/pipe21.py +++ b/pipe21.py @@ -46,6 +46,7 @@ class StarMap (B): __ror__ = lambda self, x: x | Map(lambda y: y | PipeArgs class MapApply (B): __ror__ = lambda self, it: it | Map(lambda x: x | Apply(self.f)) class Switch (B): __ror__ = lambda self, x: self.f | FilterKeys(lambda p: p(x)) | Values() | Map(lambda f: f(x)) | Pipe(next, x) class MapSwitch (B): __ror__ = lambda self, it: it | Map(lambda x: x | Switch(self.f)) +class YieldIf (B): __ror__ = lambda self, it: ((self.f or (lambda y: y))(x) for x in it if self.kw.get('key', bool)(x)) class GetItem (B): __ror__ = lambda self, x: operator.getitem(x, self.f) diff --git a/tests/pipe_test.py b/tests/pipe_test.py index 5ac9d80..4b7aea6 100644 --- a/tests/pipe_test.py +++ b/tests/pipe_test.py @@ -353,6 +353,19 @@ def test_map_switch(it, cases, expected): assert it | MapSwitch(cases) | Pipe(list) == expected +@pytest.mark.parametrize( + ('it', 'f', 'key', 'expected'), [ + (range(5), lambda x: x * 100, None, [100, 200, 300, 400]), + (range(5), lambda x: x * 100, lambda x: x % 2 == 0, [0, 200, 400]), (range(5), None, None, [1, 2, 3, 4]), + (range(5), None, lambda x: x % 2 == 0, [0, 2, 4]), + (range(5), None, None, [1, 2, 3, 4]), + ], +) +def test_yield_if(it, f, key, expected): + y = YieldIf(f) if key is None else YieldIf(f, key=key) + assert it | y | Pipe(list) == expected + + @pytest.mark.parametrize( ('it', 'f', 'expected'), [ ([('a', 1), ('b', 1), ('a', 1)], operator.add, [('a', 2), ('b', 1)]),