Map indexed

CHANGELOG.md

@@ -12,6 +12,7 @@
 
 ### New Features
 
+- Implemented `map_indexed` map 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

@@ -341,6 +341,7 @@ complete documentation reference
 | Function                                                                  | Description                                                                                                                                                                                    | Type           |
 |---------------------------------------------------------------------------|------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|----------------|
 | `map(func)/select(func)`                                                  | Maps `func` onto elements of sequence                                                                                                                                                          | transformation |
+| `map_indexed(func, start=0)`                                              | Maps `func` onto elements of sequence, with each elements index starting at `start`                                                                                                            | transformation | 
 | `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 |

functional/pipeline.py

@@ -627,6 +627,21 @@ def map(self, func: Callable[[_T_co], _T]) -> Sequence[_T]:
         """
         return self._transform(transformations.map_t(func))
 
+    def map_indexed(
+        self, func: Callable[[int, _T_co], _T], start: int = 0
+    ) -> Sequence[_T_co | _T]:
+        """
+        Maps func onto the elements of the sequence with each element's index as well.
+
+        >>> seq([1, 2, 3, 4]).map(lambda i, x: x * i)
+        [0, 2, 6, 12]
+
+        :param func: function to map with
+        :param start: beginning of index, zero by default
+        :return: sequence with func mapped onto it
+        """
+        return self._transform(transformations.map_indexed_t(func, start))
+
     def select(self, func: Callable[[_T_co], _T]) -> Sequence[_T]:
         """
         Selects f from the elements of the sequence.

functional/test/test_functional.py

@@ -348,6 +348,22 @@ def test_map(self):
         self.assertIteratorEqual(expect, result)
         self.assert_type(result)
 
+    @parametrize(
+        "sequence, start, expected",
+        [
+            ([1, 2, 3, 4], None, [0, 2, 6, 12]),
+            ([1, 2, 3, 4], 1, [1, 4, 9, 16]),
+            ([], None, []),
+        ],
+    )
+    def test_map_indexed(self, sequence, start, expected):
+        if start is None:
+            result = self.seq(sequence).map_indexed(lambda i, x: i * x)
+        else:
+            result = self.seq(sequence).map_indexed(lambda i, x: i * x, start)
+        self.assertIteratorEqual(expected, result)
+        self.assert_type(result)
+
     def test_select(self):
         l = [1, 2, 0, 5]
         expect = [2, 4, 0, 10]

functional/test/test_type.py

@@ -72,6 +72,8 @@ def type_checking() -> None:
 
         t_map: Sequence[int] = seq([1, 2, 3, 4]).map(lambda x: x * -1)
 
+        t_map_indexed: Sequence[int] = seq([1, 2, 3, 4]).map_indexed(lambda i, x: x * i)
+
         t_select: Sequence[int] = num_seq.select(lambda x: x * -1)
 
         t_starmap: Sequence[int] = seq([(2, 3), (-2, 1), (0, 10)]).starmap(

functional/transformations.py

@@ -50,6 +50,25 @@ def map_t(func: Callable):
     )
 
 
+def map_indexed_t(func: Callable, start: int):
+    """
+    Transformation for Sequence.map_indexed
+    :param func: map_indexed function
+    :param start: start of element index
+    :return: transformation
+    """
+
+    def map_indexed(map_func, sequence):
+        for i, element in enumerate(sequence, start):
+            yield map_func(i, element)
+
+    return Transformation(
+        f"map_indexed({name(func)}, {start})",
+        partial(map_indexed, func),
+        None,
+    )
+
+
 def select_t(func: Callable):
     """
     Transformation for Sequence.select