Add takewhile, dropwhile to iterators

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)
@@ -649,6 +649,52 @@ 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,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,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 "takewhile" 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 "dropwhile" 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