Add takewhile, dropwhile to iterators (JuliaLang#33437)

NEWS.md

@@ -31,6 +31,7 @@ New library functions
 * The `tempname` function now takes a `cleanup::Bool` keyword argument defaulting to `true`, which causes the process to try to ensure that any file or directory at the path returned by `tempname` is deleted upon process exit ([#33090]).
 * The `readdir` function now takes a `join::Bool` keyword argument defaulting to `false`, which when set causes `readdir` to join its directory argument with each listed name ([#33113]).
 * The new `only(x)` function returns the one-and-only element of a collection `x`, and throws an `ArgumentError` if `x` contains zero or multiple elements. ([#33129])
+* `takewhile` and `dropwhile` have been added to the Iterators submodule ([#33437]).
 
 
 Standard library changes

base/iterators.jl

@@ -22,7 +22,7 @@ import .Base:
  getindex, setindex!, get, iterate,
  popfirst!, isdone, peek
 
-export enumerate, zip, rest, countfrom, take, drop, cycle, repeated, product, flatten, partition
+export enumerate, zip, rest, countfrom, take, drop, takewhile, dropwhile, cycle, repeated, product, flatten, partition
 
 tail_if_any(::Tuple{}) = ()
 tail_if_any(x::Tuple) = tail(x)
@@ -650,6 +650,105 @@ end
 iterate(it::Drop, state) = iterate(it.xs, state)
 isdone(it::Drop, state) = isdone(it.xs, state)
 
+
+# takewhile
+
+struct TakeWhile{I,P<:Function}
+ pred::P
+ xs::I
+end
+
+"""
+ takewhile(pred, iter)
+
+An iterator that generates element from `iter` as long as predicate `pred` is true,
+afterwards, drops every element.
+
+!!! compat "Julia 1.4"
+ This function requires at least Julia 1.4.
+
+# Examples
+
+```jldoctest
+julia> s = collect(1:5)
+5-element Array{Int64,1}:
+ 1
+ 2
+ 3
+ 4
+ 5
+
+julia> collect(Iterators.takewhile(<(3),s))
+2-element Array{Int64,1}:
+ 1
+ 2
+```
+"""
+takewhile(pred,xs) = TakeWhile(pred,xs)
+
+function iterate(ibl::TakeWhile, itr...)
+ y = iterate(ibl.xs,itr...)
+ y === nothing && return nothing
+ ibl.pred(y[1]) || return nothing
+ y
+end
+
+IteratorSize(::Type{<:TakeWhile}) = SizeUnknown()
+eltype(::Type{TakeWhile{I,P}}) where {I,P} = eltype(I)
+IteratorEltype(::Type{TakeWhile{I,P}}) where {I,P} = IteratorEltype(I)
+
+
+# dropwhile
+
+struct DropWhile{I,P<:Function}
+ pred::P
+ xs::I
+end
+
+"""
+ dropwhile(pred, iter)
+
+An iterator that drops element from `iter` as long as predicate `pred` is true,
+afterwards, returns every element.
+
+!!! compat "Julia 1.4"
+ This function requires at least Julia 1.4.
+
+# Examples
+
+```jldoctest
+julia> s = collect(1:5)
+5-element Array{Int64,1}:
+ 1
+ 2
+ 3
+ 4
+ 5
+
+julia> collect(Iterators.dropwhile(<(3),s))
+3-element Array{Int64,1}:
+ 3
+ 4
+ 5
+```
+"""
+dropwhile(pred,itr) = DropWhile(pred,itr)
+
+iterate(ibl::DropWhile,itr) = iterate(ibl.xs, itr)
+function iterate(ibl::DropWhile)
+ y = iterate(ibl.xs)
+ while y !== nothing
+ ibl.pred(y[1]) || break
+ y = iterate(ibl.xs,y[2])
+ end
+ y
+end
+
+IteratorSize(::Type{<:DropWhile}) = SizeUnknown()
+eltype(::Type{DropWhile{I,P}}) where {I,P} = eltype(I)
+IteratorEltype(::Type{DropWhile{I,P}}) where {I,P} = IteratorEltype(I)
+
+
 # Cycle an iterator forever
 
 struct Cycle{I}

test/iterators.jl

@@ -149,6 +149,29 @@ for xs in Any["abc", [1, 2, 3]]
  @test take(drop(take(xs, 3), 1), 1) === take(drop(xs, 1), 1)
 end
 
+# takewhile
+# --------
+@testset begin
+ @test collect(takewhile(<(4),1:10)) == [1,2,3]
+ @test collect(takewhile(<(4),Iterators.countfrom(1))) == [1,2,3]
+ @test collect(takewhile(<(4),5:10)) == []
+ @test collect(takewhile(_->true,5:10)) == 5:10
+ @test collect(takewhile(isodd,[1,1,2,3])) == [1,1]
+ @test collect(takewhile(<(2), takewhile(<(3), [1,1,2,3]))) == [1,1]
+end
+
+# dropwhile
+# --------
+@testset begin
+ @test collect(dropwhile(<(4), 1:10)) == 4:10
+ @test collect(dropwhile(<(4), 1:10)) isa Vector{Int}
+ @test isempty(dropwhile(<(4), []))
+ @test collect(dropwhile(_->false,1:3)) == 1:3
+ @test isempty(dropwhile(_->true, 1:3))
+ @test collect(dropwhile(isodd,[1,1,2,3])) == [2,3]
+ @test collect(dropwhile(iseven,dropwhile(isodd,[1,1,2,3]))) == [3]
+end
+
 # cycle
 # -----
 let i = 0