adds the nth function for iterables

base/iterators.jl

@@ -16,7 +16,7 @@ using .Base:
     (:), |, +, -, *, !==, !, ==, !=, <=, <, >, >=, =>, missing,
     any, _counttuple, eachindex, ntuple, zero, prod, reduce, in, firstindex, lastindex,
     tail, fieldtypes, min, max, minimum, zero, oneunit, promote, promote_shape, LazyString,
-    afoldl
+    afoldl, mod1
 using .Core
 using Core: @doc
 
@@ -32,7 +32,7 @@ import Base:
     getindex, setindex!, get, iterate,
     popfirst!, isdone, peek, intersect
 
-export enumerate, zip, rest, countfrom, take, drop, takewhile, dropwhile, cycle, repeated, product, flatten, flatmap, partition
+export enumerate, zip, rest, countfrom, take, drop, takewhile, dropwhile, cycle, repeated, product, flatten, flatmap, partition, nth
 public accumulate, filter, map, peel, reverse, Stateful
 
 """
@@ -1602,4 +1602,81 @@ end
 # be the same as the keys, so this is a valid optimization (see #51631)
 pairs(s::AbstractString) = IterableStatePairs(s)
 
+"""
+    nth(itr, n::Integer)
+
+Get the `n`th element of an iterable collection. Throw a `BoundsError`[@ref] if not existing.
+Will advance any `Stateful`[@ref] iterator.
+
+See also: [`first`](@ref), [`last`](@ref), [`nth`](@ref)
+
+# Examples
+```jldoctest
+julia> Iterators.nth(2:2:10, 4)
+8
+
+julia> Iterators.nth(reshape(1:30, (5,6)), 6)
+6
+
+julia> stateful = Iterators.Stateful(1:10); nth(stateful, 7)
+7
+
+julia> first(stateful)
+8
+```
+"""
+nth(itr, n::Integer) = _nth(itr, n)
+
+# Count
+nth(itr::Count, n::Integer) = n > 0 ? itr.start + itr.step * (n - 1) : throw(ArgumentError("n must be positive."))
+# Repeated
+nth(itr::Repeated, ::Integer) = itr.x
+# Take(Repeated)
+nth(itr::Take{<:Repeated}, n::Integer) =
+    n > itr.n ? throw(BoundsError(itr, n)) : nth(itr.xs, n)
+
+# infinite cycle
+function nth(itr::Cycle{I}, n::Integer) where {I}
+    if IteratorSize(I) isa Union{HasShape, HasLength}
+        _nth(itr.xs, mod1(n, length(itr.xs)))
+    else
+        _nth(itr, n)
+    end
+end
+
+# finite cycle: in reality a Flatten{Take{Repeated{O}}} iterator
+function nth(itr::Flatten{Take{Repeated{O}}}, n::Integer) where {O}
+    if IteratorSize(O) isa Union{HasShape, HasLength}
+        cycles = itr.it.n
+        repeated = itr.it.xs.x
+        k = length(repeated)
+        n > k*cycles ? throw(BoundsError(itr, n)) : _nth(repeated, mod1(n, k))
+    else
+        _nth(itr, n)
+    end
+end
+
+_nth(itr, n) = first(drop(itr, n-1))
+_nth(itr::AbstractArray, n) = itr[begin + n-1]
+"""
+    nth(n::Integer)
+
+Return a function that gets the `n`-th element from any iterator passed to it.
+Equivalent to `Base.Fix2(nth, n)` or `itr -> nth(itr, n)`.
+
+See also: [`nth`](@ref), [`Base.Fix2`](@ref)
+# Examples
+```jldoctest
+julia> fifth_element = Iterators.nth(5)
+(::Base.Fix2{typeof(Base.Iterators.nth), Int64}) (generic function with 2 methods)
+
+julia> fifth_element(reshape(1:30, (5,6)))
+5
+
+julia> map(fifth_element, ("Willis", "Jovovich", "Oldman"))
+('i', 'v', 'a')
+```
+"""
+nth(n::Integer) = Base.Fix2(nth, n)
+
 end

test/iterators.jl

@@ -1133,6 +1133,30 @@ end
     end
 end
 
+@testset "nth" begin
+    Z = Array{Int,0}(undef)
+    Z[] = 17
+    itrs = Any[collect(1:1000), 10:6:1000, "∀ϵ>0", (1, 3, 5, 10, 78), reshape(1:30, (5, 6)),
+        Z, 3, true, 'x', 4=>5, view(Z), view(reshape(1:30, (5, 6)), 2:4, 2:6),
+        (x^2 for x in 1:10), Iterators.Filter(isodd, 1:10), Iterators.flatten((1:10, 50:60)),
+        pairs(50:60), zip(1:10, 21:30, 51:60), Iterators.product(1:3, 10:12),
+        Iterators.repeated(3.14159, 5), (a=2, b=3, c=5, d=7, e=11), Iterators.cycle(collect(1:100)),
+        Iterators.cycle([1, 2, 3, 4, 5], 5)]
+    ns = Any[
+        234, 123, 3, 2, 21, 1, 1, 1, 1, 1, 1, 10, 9, 3, 15, 7, 6, 3, 4, 4, 99999, 25
+    ]
+    expected = Any[
+        234, 742, '>', 3, 21, 17, 3, true, 'x', 4, 17, 22, 81, 5, 54, (7=>56), (6, 26, 56), (3, 10), 3.14159, 7, 99, 5
+    ]
+    @test length(itrs) == length(ns) == length(expected)
+    testset = zip(itrs, ns, expected)
+    @testset "iter: $IT" for (IT, n, exp) in testset
+        @test exp == nth(IT, n)
+        IT isa Iterators.Cycle && continue # cycles are infinite so never OOB
+        @test_throws Union{ArgumentError,BoundsError} nth(IT, 999999999)
+    end
+end
+
 @testset "Iterators docstrings" begin
     @test isempty(Docs.undocumented_names(Iterators))
 end