adds the nth function for iterables

base/iterators.jl

@@ -30,7 +30,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
 
 """
@@ -1595,4 +1595,102 @@ 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> nth(2:2:10, 4)
+8
+
+julia> nth(reshape(1:30, (5,6)), 6)
+6
+
+julia> stateful = Stateful(1:10); nth(stateful, 7)
+7
+
+julia> first(stateful)
+8
+```
+"""
+nth(itr, n::Integer) = _nth(IteratorSize(itr), itr, n)
+nth(itr::AbstractArray, n::Integer) = _withlength_nth_throw(itr, n, length(itr))
+# specialized versions to better interact with existing iterators
+# Count
+nth(itr::Count, n::Integer) = n > 0 ? itr.start + itr.step * (n - 1) : throw(ArgumentError("n must be positive."))
+# # Repeated
+nth(itr::Repeated, n::Integer) = itr.x
+# # Take(Repeated)
+nth(itr::Take{Repeated{O}}, n::Integer) where {O} = begin
+    n > itr.n ? throw(BoundsError("attempted to access $(itr.n)-element $(typeof(itr)) at position $n")) : itr.xs.x
+end
+# infinite cycle
+nth(itr::Cycle{I}, n::Integer) where {I} = _nth_inf_cycle(IteratorSize(I), itr, n)
+# finite cycle: in reality a Flatten{Take{Repeated{O}}} iterator
+nth(itr::Flatten{Take{Repeated{O}}}, n::Integer) where {O} = _nth_finite_cycle(IteratorSize(O), itr, n)
+
+_nth(::SizeUnknown, itr, n) = _fallback_nth_throw(itr, n)
+_nth(::Union{HasShape,HasLength}, itr, n) = _withlength_nth_throw(itr, n, length(itr))
+_nth(::IsInfinite, itr, n) = _inbounds_nth(itr, n)
+
+_inbounds_nth(itr, n) = iterate(drop(itr, n - 1))[1]
+_inbounds_nth(itr::AbstractArray, n) = itr[begin + n-1]
+
+function _withlength_nth_throw(itr, n, N)
+    n > N && throw(BoundsError("attempted to access $N-element $(typeof(itr)) at position $n"))
+    _inbounds_nth(itr, n)
+end
+
+function _fallback_nth_throw(itr, n)
+    y = iterate(drop(itr, n - 1))
+    y === nothing && throw(BoundsError("Iterator $(typeof(itr)) has less than $n elements"))
+    y[1]
+end
+
+_nth_inf_cycle(::IsInfinite, itr, n) = _inbounds_nth(itr.xs, n)
+_nth_inf_cycle(::SizeUnknown, itr, n) = _fallback_nth(itr.xs, n)
+_nth_inf_cycle(::Union{HasShape,HasLength}, itr, n) = _repeating_cycle_nth(itr.xs, n, length(itr.xs))
+
+_nth_finite_cycle(::IsInfinite, itr, n) = _inbounds_nth(itr, n)
+_nth_finite_cycle(::SizeUnknown, itr, n) = _fallback_nth_throw(itr, n)
+_nth_finite_cycle(::Union{HasShape,HasLength}, itr, n) = _walk_cycle_throw(itr, n)
+
+_repeating_cycle_nth(inner_itr, n, inner_N) = _inbounds_nth(inner_itr, 1 + ((n - 1) % inner_N))
+
+function _walk_cycle_throw(itr, n)
+    N = itr.it.n # `Take` iterator n
+    torepeat = itr.it.xs.x # repeated object
+    K = length(torepeat)
+    n > K * N && throw(BoundsError("attempted to access $(N*K)-element $(typeof(itr)) at position $n"))
+    _repeating_cycle_nth(torepeat, n, K)
+end
+
+"""
+    nth(n::Integer)
+
+Return a function that gets the `n`-th element from any iterator passed to it.
+Throw a `BoundsError`[@ref] if not existing.
+
+Fixes the second element. Equivalent to `Base.Fix2(nth, n)`.
+Will advance any `Stateful`[@ref] iterator.
+
+See also: [`nth`](@ref), [`Base.Fix2`](@ref)
+# Examples
+```jldoctest
+julia> fifth_element = nth(5)
+(::Base.Fix2{typeof(nth), Int64}) (generic function with 1 method)
+
+julia> fifth_element(reshape(1:30, (5,6)))
+5
+
+julia> map(nth(3), my_vec)
+```
+"""
+nth(n::Integer) = Base.Fix2(nth, n)
+
 end

test/iterators.jl

@@ -1074,6 +1074,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 Cycle && continue # cycles are infinite so never OOB
+        @test nth(IT, 999999999) === nothing
+    end
+end
+
 @testset "Iterators docstrings" begin
     @test isempty(Docs.undocumented_names(Iterators))
 end