Filter indexed

CHANGELOG.md

@@ -12,6 +12,7 @@
 
 ### New Features
 
+- Implemented `filter_indexed` filter function that provides both the index of the element and the element as arguments  
 - Added `first_or_none`, a function to match `head_or_none`   
 - Added run_test.sh script
 - Added [parametrize](https://pypi.org/project/parametrize/) for parameterized unit tests

README.md

@@ -344,6 +344,7 @@ complete documentation reference
 | `starmap(func)/smap(func)`                                                | Applies `func` to sequence with `itertools.starmap`                                                                                                                                            | transformation |
 | `filter(func)/where(func)`                                                | Filters elements of sequence to only those where `func(element)` is `True`                                                                                                                     | transformation |
 | `filter_not(func)`                                                        | Filters elements of sequence to only those where `func(element)` is `False`                                                                                                                    | transformation |
+| `filter_indexed(func, start=0)`                                           | Filters elements of sequence to only those where `func(index, element)` is `True`, with each element's index starting at `start`                                                               | transformation |
 | `flatten()`                                                               | Flattens sequence of lists to a single sequence                                                                                                                                                | transformation |
 | `flat_map(func)`                                                          | Maps `func` to each element, then merges the result to one flat sequence. `func` must return an iterable                                                                                       | transformation |
 | `group_by(func)`                                                          | Groups sequence into `(key, value)` pairs where `key=func(element)` and `value` is from the original sequence                                                                                  | transformation |

functional/pipeline.py

@@ -803,6 +803,22 @@ def filter_not(self, func: Callable[[_T_co], Any]) -> Sequence[_T_co]:
         """
         return self._transform(transformations.filter_not_t(func))
 
+    def filter_indexed(
+        self, func: Callable[[int, _T], bool], start: int = 0
+    ) -> Sequence[_T]:
+        """
+        Filters sequence with predicate func that also provides each element's index,
+        which can be used to filter on.
+
+        >>> seq([-1, 1, -2, 2]).filter_indexed(lambda i, x: x < 0 and i > 1)
+        [-2]
+
+        :param func: function to filter on
+        :param start: beginning of index, zero by default
+        :return: filtered sequence
+        """
+        return self._transform(transformations.filter_indexed_t(func, start))
+
     def where(self, func: Callable[[_T_co], Any]) -> Sequence[_T_co]:
         """
         Selects elements where func evaluates to True.

functional/test/test_functional.py

@@ -382,6 +382,24 @@ def test_filter(self):
         self.assertIteratorEqual(expect, result)
         self.assert_type(result)
 
+    @parametrize(
+        "sequence, start, expected",
+        [
+            ([1, 3, 5, 7, 9, 11], None, [5, 7, 9]),
+            ([1, 3, 5, 7, 9, 11], 1, [3, 5, 7, 9]),
+            ([], None, []),
+        ],
+    )
+    def test_filter_indexed(self, sequence, start, expected):
+        if start is None:
+            result = self.seq(sequence).filter_indexed(lambda i, x: i > 1 and x < 10)
+        else:
+            result = self.seq(sequence).filter_indexed(
+                lambda i, x: i > 1 and x < 10, start
+            )
+        self.assertIteratorEqual(expected, result)
+        self.assert_type(result)
+
     def test_where(self):
         def f(x):
             return x > 0

functional/test/test_type.py

@@ -92,6 +92,10 @@ def type_checking() -> None:
 
         t_filter_not: Sequence[int] = seq([-1, 1, -2, 2]).filter_not(lambda x: x > 0)
 
+        t_filter_indexed: Sequence[int] = seq([-1, 1, -2, 2]).filter_indexed(
+            lambda i, x: x > 0 and i > 1
+        )
+
         t_where: Sequence[int] = seq([-1, 1, -2, 2]).where(lambda x: x > 0)
 
         t_count: int = seq([-1, -2, 1, 2]).count(lambda x: x > 0)

functional/transformations.py

@@ -115,6 +115,24 @@ def filter_not_t(func: Callable):
     )
 
 
+def filter_indexed_t(func: Callable, start: int):
+    """
+    Transformation for Sequence.filter_indexed
+    :param func: filter_indexed function
+    :param start: start of element index
+    :return: transformation
+    """
+
+    def filter_indexed(predicate, sequence):
+        for i, element in enumerate(sequence, start):
+            if predicate(i, element):
+                yield element
+
+    return Transformation(
+        f"filter_indexed({name(func)}, {start})", partial(filter_indexed, func), None
+    )
+
+
 def reversed_t():
     """
     Transformation for Sequence.reverse