Merge pull request #66 from JuliaCollections/rf/firstrest

Add a firstrest function

docs/src/index.md

@@ -20,6 +20,13 @@ Iterate through values skipping over those already encountered.
 distinct
 ```
 
+## firstrest(xs)
+
+Return first element and `Iterators.rest` iterator as a tuple.
+```@docs
+firstrest
+```
+
 ## groupby(f, xs)
 
 Group consecutive values that share the same result of applying `f`.

src/IterTools.jl

@@ -9,6 +9,7 @@ import Base: SizeUnknown, IsInfinite, HasLength, HasShape
 import Base: HasEltype, EltypeUnknown
 
 export
+    firstrest,
     takestrict,
     repeatedly,
     chain,
@@ -84,6 +85,29 @@ macro ifsomething(ex)
     end
 end
 
+"""
+    firstrest(xs) -> (f, r)
+
+Return the first element and an iterator of the rest as a tuple.
+
+```jldoctest
+julia> f, r = firstrest(1:3)
+(1, Base.Iterators.Rest{UnitRange{Int64},Int64}(1:3, 1))
+
+julia> collect(r)
+2-element Array{Int64,1}:
+ 2
+ 3
+```
+"""
+function firstrest(xs)
+    t = iterate(xs)
+    t === nothing && throw(ArgumentError("collection must be non-empty"))
+    f, s = t
+    r = Iterators.rest(xs, s)
+    return f, r
+end
+
 # Iterate through the first n elements, throwing an exception if
 # fewer than n items ar encountered.
 

test/runtests.jl

@@ -10,6 +10,36 @@ include("testing_macros.jl")
 
 @testset "IterTools" begin
 @testset "iterators" begin
+    @testset "firstrest" begin
+        # Ranges/generators have different rest states vs array/tuples
+        test_base_cases = [
+            (1:1, 1),
+            (1:3, 1),
+            ([1], 2),
+            ([1, 2, 3], 2),
+            ((1,), 2),
+            ((1, 2, 3), 2),
+            ((i for i in 1:1), 1),
+            ((i for i in 1:3), 1),
+        ]
+        @testset "$xs" for (xs, s) in test_base_cases
+            f, r = firstrest(xs)
+            @test f == first(xs)
+            @test collect(r) == collect(Iterators.rest(xs, s))
+        end
+
+        test_empty_cases = [
+            (1:0, 1),
+            (Int[], 2),
+            ((), 2),
+            ((i for i in 1:0), 1),
+        ]
+
+        @testset "$xs" for (xs, s) in test_empty_cases
+            @test_throws ArgumentError firstrest(xs)
+        end
+    end
+
     @testset "takestrict" begin
         itr = 1:10
         take_itr = takestrict(itr, 5)