add more commands to std iter (#9129)

# Description

this pr adds the following commands to `std iter`
- `iter find-index` -> returns the index of the first element that
matches the predicate or `-1` if none
- `iter flat-map` -> maps a closure to each nested structure and
flattens the result
- `iter zip-with` -> zips two structures and applies a closure to each
of the zips

it also fixes some \*\*very minor\*\* inconsistencies in the module
This commit is contained in:
mike 2023-05-10 20:18:42 +03:00 committed by GitHub
parent 43a3983d36
commit 39e51f1953
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 139 additions and 10 deletions

View file

@ -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 # Returns a new list with the separator between adjacent
# items of the original list # 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] # assert equal $res [1 0 2 0 3 0 4]
# ``` # ```
export def "iter intersperse" [ # -> list<any> export def "iter intersperse" [ # -> list<any>
separator: any, # the separator to be used separator: any # the separator to be used
] { ] {
reduce -f [] {|it, acc| reduce -f [] {|it, acc|
$acc ++ [$it, $separator] $acc ++ [$it, $separator]
} }
| match $in { | match $in {
[] => [], [] => [],
@ -83,14 +122,14 @@ export def "iter scan" [ # -> list<any>
fn: closure # the closure to perform the scan fn: closure # the closure to perform the scan
--noinit(-n) # remove the initial value from the result --noinit(-n) # remove the initial value from the result
] { ] {
reduce -f [$init] {|it, acc| reduce -f [$init] {|it, acc|
$acc ++ [(do $fn ($acc | last) $it)] $acc ++ [(do $fn ($acc | last) $it)]
} }
| if $noinit { | if $noinit {
$in | skip $in | skip
} else { } else {
$in $in
} }
} }
# Returns a list of values for which the supplied closure does not # Returns a list of values for which the supplied closure does not
@ -119,3 +158,42 @@ export def "iter filter-map" [ # -> list<any>
$it != null $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<any>
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<any>
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 }
}
}

View file

@ -59,3 +59,54 @@ export def test_iter_filter_map [] {
) )
assert equal $res [3 42 69] 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]
}