diff --git a/NEWS.md b/NEWS.md index 1e4e24b67a4db..55b5a09d5b525 100644 --- a/NEWS.md +++ b/NEWS.md @@ -5,6 +5,7 @@ New language features --------------------- - New `Base.@acquire` macro for a non-closure version of `Base.acquire(f, s::Base.Semaphore)`, like `@lock`. ([#56845]) + - New `nth` function to access the `n`-th element of a generic iterable. ([#56580]) - The character U+1F8B2 🢲 (RIGHTWARDS ARROW WITH LOWER HOOK), newly added by Unicode 16, is now a valid operator with arrow precedence, accessible as `\hookunderrightarrow` at the REPL. ([JuliaLang/JuliaSyntax.jl#525], [#57143]) diff --git a/base/iterators.jl b/base/iterators.jl index 4b956073f0c04..53d7b28316ee1 100644 --- a/base/iterators.jl +++ b/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 """ @@ -991,6 +991,7 @@ end reverse(it::Cycle) = Cycle(reverse(it.xs)) last(it::Cycle) = last(it.xs) + # Repeated - repeat an object infinitely many times struct Repeated{O} @@ -1602,4 +1603,83 @@ 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) + +# 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); Iterators.nth(stateful, 7) +7 + +julia> first(stateful) +8 +``` +""" +nth(itr, n::Integer) = _nth(IteratorSize(itr), itr, n) +nth(itr::Cycle{I}, n::Integer) where I = _nth(IteratorSize(I), itr, n) +nth(itr::Flatten{Take{Repeated{O}}}, n::Integer) where O = _nth(IteratorSize(O), itr, n) +@propagate_inbounds nth(itr::AbstractArray, n::Integer) = itr[begin + n - 1] + +function _nth(::Union{HasShape, HasLength}, itr::Cycle{I}, n::Integer) where {I} + N = length(itr.xs) + N == 0 && throw(BoundsError(itr, n)) + + # prevents wrap around behaviour and inherit the error handling + return nth(itr.xs, n > 0 ? mod1(n, N) : n) +end + +# Flatten{Take{Repeated{O}}} is the actual type of an Iterators.cycle(iterable::O, m) iterator +function _nth(::Union{HasShape, HasLength}, itr::Flatten{Take{Repeated{O}}}, n::Integer) where {O} + cycles = itr.it.n + torepeat = itr.it.xs.x + k = length(torepeat) + (n > k*cycles || k == 0) && throw(BoundsError(itr, n)) + + # prevent wrap around behaviour and inherit the error handling + return nth(torepeat, n > 0 ? mod1(n, k) : n) +end + +function _nth(::IteratorSize, itr, n::Integer) + # unrolled version of `first(drop)` + n > 0 || throw(BoundsError(itr, n)) + y = iterate(itr) + for _ in 1:n-1 + y === nothing && break + y = iterate(itr, y[2]) + end + y === nothing && throw(BoundsError(itr, n)) + y[1] +end +""" + 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 diff --git a/test/iterators.jl b/test/iterators.jl index a6ab4720c0d0c..cebabf9cc07fc 100644 --- a/test/iterators.jl +++ b/test/iterators.jl @@ -1139,6 +1139,65 @@ end end end +@testset "nth" begin + + Z = Array{Int,0}(undef) + Z[] = 17 + it_result_pairs = Dict( + (Z, 1) => 17, + (collect(1:100), 23) => 23, + (10:6:1000, 123) => 10 + 6 * 122, + ("∀ϵ>0", 3) => '>', + ((1, 3, 5, 10, 78), 2) => 3, + (reshape(1:30, (5, 6)), 21) => 21, + (3, 1) => 3, + (true, 1) => true, + ('x', 1) => 'x', + (4 => 5, 2) => 5, + (view(Z), 1) => 17, + (view(reshape(1:30, (5, 6)), 2:4, 2:6), 10) => 22, + ((x^2 for x in 1:10), 9) => 81, + (Iterators.Filter(isodd, 1:10), 3) => 5, + (Iterators.flatten((1:10, 50:60)), 15) => 54, + (pairs(50:60), 7) => 7 => 56, + (zip(1:10, 21:30, 51:60), 6) => (6, 26, 56), + (Iterators.product(1:3, 10:12), 3) => (3, 10), + (Iterators.repeated(3.14159, 5), 4) => 3.14159, + ((a=2, b=3, c=5, d=7, e=11), 4) => 7, + (Iterators.cycle(collect(1:100)), 9999) => 99, + (Iterators.cycle([1, 2, 3, 4, 5], 5), 25) => 5, + (Iterators.cycle("String", 10), 16) => 'i', + (Iterators.cycle(((),)), 1000) => () + ) + + + @testset "iter: $IT" for (IT, n) in keys(it_result_pairs) + @test it_result_pairs[(IT, n)] == nth(IT, n) + @test_throws BoundsError nth(IT, -42) + + IT isa Iterators.Cycle && continue # cycles are infinite so never OOB + @test_throws BoundsError nth(IT, 999999999) + end + + empty_cycle = Iterators.cycle([]) + @test_throws BoundsError nth(empty_cycle, 42) + + # test the size unknown branch for cycles + # only generate odd numbers so we know the actual length + # but the iterator is still SizeUnknown() + it_size_unknown = Iterators.filter(isodd, 1:2:10) + @test Base.IteratorSize(it_size_unknown) isa Base.SizeUnknown + @test length(collect(it_size_unknown)) == 5 + + cycle_size_unknown = Iterators.cycle(it_size_unknown) + finite_cycle_size_unknown = Iterators.cycle(it_size_unknown, 5) + @test nth(cycle_size_unknown, 2) == 3 + @test nth(cycle_size_unknown, 20) == 9 # mod1(20, 5) = 5, wraps 4 times + @test nth(finite_cycle_size_unknown, 2) == 3 + @test nth(finite_cycle_size_unknown, 20) == 9 + @test_throws BoundsError nth(finite_cycle_size_unknown, 30) # only wraps 5 times, max n is 5 * 5 = 25 +end + @testset "Iterators docstrings" begin @test isempty(Docs.undocumented_names(Iterators)) end