From 3c9613837fc90c1ccebd5b3fc0da5af7be7ef472 Mon Sep 17 00:00:00 2001 From: Bhargav Andhe Date: Tue, 18 Aug 2026 22:42:46 +0530 Subject: [PATCH] feat: add Seq.try_find_index --- expression/collections/seq.py | 33 +++++++++++++++++++++++ tests/test_seq.py | 51 +++++++++++++++++++++++++++++++++++ 2 files changed, 84 insertions(+) diff --git a/expression/collections/seq.py b/expression/collections/seq.py index 5ece6d2..f5c1c8a 100644 --- a/expression/collections/seq.py +++ b/expression/collections/seq.py @@ -31,8 +31,10 @@ from typing import TYPE_CHECKING, Any, TypeVar, cast, overload from expression.core import ( + Nothing, Option, PipeMixin, + Some, SupportsGreaterThan, SupportsLessThan, SupportsSum, @@ -322,6 +324,10 @@ def take(self, count: int) -> Seq[_TSource]: def to_list(self) -> Block[_TSource]: return to_list(self) + def try_find_index(self, predicate: Callable[[_TSource], bool]) -> Option[int]: + """Return the index of the first element matching the predicate, if any.""" + return pipe(self, try_find_index(predicate)) + def dict(self) -> Iterable[_TSource]: """Returns a json serializable representation of the list.""" @@ -942,6 +948,32 @@ def to_list(source: Iterable[_TSource]) -> Block[_TSource]: return Block.of_seq(source) +@curry_flip(1) +def try_find_index(source: Iterable[_TSource], predicate: Callable[[_TSource], bool]) -> Option[int]: + """Return the index of the first element matching the predicate, if any. + + Indices are zero-based, and evaluation stops as soon as the predicate + returns `True`. + + Args: + source: The input sequence. + predicate: A function to test each element. + + Returns: + The first matching index wrapped in `Some`, or `Nothing` when no + element matches. + + Example: + >>> pipe([1, 2, 3], try_find_index(lambda value: value % 2 == 0)) + Some 1 + """ + for index, value in enumerate(source): + if predicate(value): + return Some(index) + + return Nothing + + @curry_flip(1) def unfold(state: _TState, generator: Callable[[_TState], Option[tuple[_TSource, _TState]]]) -> Iterable[_TSource]: """Unfold sequence. @@ -1035,6 +1067,7 @@ def _zip(source2: Iterable[_TResult]) -> Iterable[tuple[_TSource, _TResult]]: "sum_by", "tail", "take", + "try_find_index", "unfold", "zip", ] diff --git a/tests/test_seq.py b/tests/test_seq.py index 959f673..a013e55 100644 --- a/tests/test_seq.py +++ b/tests/test_seq.py @@ -363,6 +363,57 @@ def source() -> Iterable[int]: assert consumed == [1, 2] +def test_seq_try_find_index_pipe_returns_first_matching_index(): + is_even: Callable[[int], bool] = lambda value: value % 2 == 0 + result = pipe([1, 3, 4, 6], seq.try_find_index(is_even)) + + assert result == Some(2) + + +def test_seq_try_find_index_fluent(): + source = Seq[int].of_iterable([1, 3, 4, 6]) + + assert source.try_find_index(lambda value: value % 2 == 0) == Some(2) + + +def test_seq_try_find_index_can_return_zero(): + is_even: Callable[[int], bool] = lambda value: value % 2 == 0 + + assert pipe([2, 4], seq.try_find_index(is_even)) == Some(0) + + +def test_seq_try_find_index_returns_nothing_without_match(): + empty: list[int] = [] + always_true: Callable[[int], bool] = lambda _: True + is_even: Callable[[int], bool] = lambda value: value % 2 == 0 + + assert pipe(empty, seq.try_find_index(always_true)) is Nothing + assert pipe([1, 3], seq.try_find_index(is_even)) is Nothing + + +def test_seq_try_find_index_stops_after_first_match(): + consumed: list[int] = [] + is_even: Callable[[int], bool] = lambda value: value % 2 == 0 + + def source() -> Iterable[int]: + for value in [1, 3, 4, 6]: + consumed.append(value) + yield value + + result = pipe(source(), seq.try_find_index(is_even)) + + assert result == Some(2) + assert consumed == [1, 3, 4] + + +def test_seq_try_find_index_propagates_predicate_exceptions(): + def predicate(_: int) -> bool: + raise ValueError("predicate failed") + + with pytest.raises(ValueError, match="predicate failed"): + pipe([1], seq.try_find_index(predicate)) + + rtn: Callable[[int], Seq[int]] = seq.singleton empty: Seq[int] = seq.empty