diff --git a/crates/nu-std/lib/iter.nu b/crates/nu-std/lib/iter.nu index 24a0952056..8efe636506 100644 --- a/crates/nu-std/lib/iter.nu +++ b/crates/nu-std/lib/iter.nu @@ -37,6 +37,45 @@ export def "iter find" [ # -> any | null } } +# Returns the index of the first element that matches the predicate or +# -1 if none +# +# # Invariant +# > The closure has to return a bool +# +# # Example +# ```nu +# use std ["assert equal" "iter find-index"] +# +# let res = ( +# ["iter", "abc", "shell", "around", "nushell", "std"] +# | iter find-index {|x| $x starts-with 's'} +# ) +# assert equal $res 2 +# +# let is_even = {|x| $x mod 2 == 0} +# let res = ([3 5 13 91] | iter find-index $is_even) +# assert equal $res -1 +# ``` +export def "iter find-index" [ # -> int + fn: closure # the closure used to perform the search +] { + let matches = ( + enumerate + | each {|it| + if (do $fn $it.item) { + $it.index + } + } + ) + + if ($matches | is-empty) { + -1 + } else { + $matches | first + } +} + # Returns a new list with the separator between adjacent # items of the original list # @@ -48,10 +87,10 @@ export def "iter find" [ # -> any | null # assert equal $res [1 0 2 0 3 0 4] # ``` export def "iter intersperse" [ # -> list - separator: any, # the separator to be used + separator: any # the separator to be used ] { reduce -f [] {|it, acc| - $acc ++ [$it, $separator] + $acc ++ [$it, $separator] } | match $in { [] => [], @@ -83,14 +122,14 @@ export def "iter scan" [ # -> list fn: closure # the closure to perform the scan --noinit(-n) # remove the initial value from the result ] { - reduce -f [$init] {|it, acc| - $acc ++ [(do $fn ($acc | last) $it)] - } - | if $noinit { - $in | skip - } else { - $in - } + reduce -f [$init] {|it, acc| + $acc ++ [(do $fn ($acc | last) $it)] + } + | if $noinit { + $in | skip + } else { + $in + } } # Returns a list of values for which the supplied closure does not @@ -119,3 +158,42 @@ export def "iter filter-map" [ # -> list $it != null } } + +# Maps a closure to each nested structure and flattens the result +# +# # Example +# ```nu +# use std ["assert equal" "iter flat-map"] +# +# let res = ( +# [[1 2 3] [2 3 4] [5 6 7]] | iter flat-map {|it| $it | math sum} +# ) +# assert equal $res [6 9 18] +# ``` +export def "iter flat-map" [ # -> list + fn: closure # the closure to map to the nested structures +] { + each {|it| do $fn $it } | flatten +} + +# Zips two structures and applies a closure to each of the zips +# +# # Example +# ```nu +# use std ["assert equal" "iter iter zip-with"] +# +# let res = ( +# [1 2 3] | iter zip-with [2 3 4] {|a, b| $a + $b } +# ) +# +# assert equal $res [3 5 7] +# ``` +export def "iter zip-with" [ # -> list + other: any # the structure to zip with + fn: closure # the closure to apply to the zips +] { + zip $other + | each {|it| + reduce {|it, acc| do $fn $acc $it } + } +} diff --git a/crates/nu-std/tests/test_iter.nu b/crates/nu-std/tests/test_iter.nu index 10cb1597f7..52f7bff14e 100644 --- a/crates/nu-std/tests/test_iter.nu +++ b/crates/nu-std/tests/test_iter.nu @@ -59,3 +59,54 @@ export def test_iter_filter_map [] { ) assert equal $res [3 42 69] } + +export def test_iter_find_index [] { + let res = ( + ["iter", "abc", "shell", "around", "nushell", "std"] + | iter find-index {|x| $x starts-with 's'} + ) + assert equal $res 2 + + let is_even = {|x| $x mod 2 == 0} + let res = ([3 5 13 91] | iter find-index $is_even) + assert equal $res (-1) + + let res = (42 | iter find-index {|x| $x == 42}) + assert equal $res 0 +} + +export def test_iter_zip_with [] { + let res = ( + [1 2 3] | iter zip-with [2 3 4] {|a, b| $a + $b } + ) + + assert equal $res [3 5 7] + + let res = (42 | iter zip-with [1 2 3] {|a, b| $a // $b}) + assert equal $res [42] + + let res = (2..5 | iter zip-with 4 {|a, b| $a * $b}) + assert equal $res [8] + + let res = ( + [[name repo]; [rust github] [haskell gitlab]] + | iter zip-with 1.. {|data, num| + { name: $data.name, repo: $data.repo position: $num } + } + ) + assert equal $res [ + [name repo position]; + [rust github 1] + [haskell gitlab 2] + ] +} + +export def test_iter_flat_map [] { + let res = ( + [[1 2 3] [2 3 4] [5 6 7]] | iter flat-map {|it| $it | math sum} + ) + assert equal $res [6 9 18] + + let res = ([1 2 3] | iter flat-map {|it| $it + ($it * 10)}) + assert equal $res [11 22 33] +}