diff --git a/NEWS.md b/NEWS.md index d0c2437ed5c0a..dd86f5c8409c1 100644 --- a/NEWS.md +++ b/NEWS.md @@ -65,6 +65,8 @@ Standard library changes * `unique(f, itr; seen=Set{T}())` now allows you to declare the container type used for keeping track of values returned by `f` on elements of `itr` ([#36280]). * `Libdl` has been moved to `Base.Libc.Libdl`, however it is still accessible as an stdlib ([#35628]). +* `first` and `last` functions now accept an integer as second argument to get that many + leading or trailing elements of any iterable ([#34868]). #### LinearAlgebra * New method `LinearAlgebra.issuccess(::CholeskyPivoted)` for checking whether pivoted Cholesky factorization was successful ([#36002]). diff --git a/base/abstractarray.jl b/base/abstractarray.jl index 504b51efec3e5..68dce7e7d0b20 100644 --- a/base/abstractarray.jl +++ b/base/abstractarray.jl @@ -343,6 +343,33 @@ function first(itr) x[1] end +""" + first(itr, n::Integer) + +Get the first `n` elements of the iterable collection `itr`, or fewer elements if `v` is not +long enough. + +# Examples +```jldoctest +julia> first(["foo", "bar", "qux"], 2) +2-element Vector{String}: + "foo" + "bar" + +julia> first(1:6, 10) +1:6 + +julia> first(Bool[], 1) +Bool[] +``` +""" +first(itr, n::Integer) = collect(Iterators.take(itr, n)) +# Faster method for vectors +function first(v::AbstractVector, n::Integer) + n < 0 && throw(ArgumentError("Number of elements must be nonnegative")) + @inbounds v[begin:min(begin + n - 1, end)] +end + """ last(coll) @@ -361,6 +388,33 @@ julia> last([1; 2; 3; 4]) """ last(a) = a[end] +""" + last(itr, n::Integer) + +Get the last `n` elements of the iterable collection `itr`, or fewer elements if `v` is not +long enough. + +# Examples +```jldoctest +julia> last(["foo", "bar", "qux"], 2) +2-element Vector{String}: + "bar" + "qux" + +julia> last(1:6, 10) +1:6 + +julia> last(Float64[], 1) +Float64[] +``` +""" +last(itr, n::Integer) = reverse!(collect(Iterators.take(Iterators.reverse(itr), n))) +# Faster method for arrays +function last(v::AbstractArray, n::Integer) + n < 0 && throw(ArgumentError("Number of elements must be nonnegative")) + @inbounds v[max(begin, end - n + 1):end] +end + """ strides(A) diff --git a/test/abstractarray.jl b/test/abstractarray.jl index 335697b65f8eb..2db1638c8c950 100644 --- a/test/abstractarray.jl +++ b/test/abstractarray.jl @@ -1126,3 +1126,18 @@ end end end end + +@testset "first/last n elements of $(typeof(itr))" for itr in (collect(1:9), + [1 4 7; 2 5 8; 3 6 9], + ntuple(identity, 9)) + @test first(itr, 6) == [itr[1:6]...] + @test first(itr, 25) == [itr[:]...] + @test first(itr, 25) !== itr + @test first(itr, 1) == [itr[1]] + @test_throws ArgumentError first(itr, -6) + @test last(itr, 6) == [itr[end-5:end]...] + @test last(itr, 25) == [itr[:]...] + @test last(itr, 25) !== itr + @test last(itr, 1) == [itr[end]] + @test_throws ArgumentError last(itr, -6) +end diff --git a/test/offsetarray.jl b/test/offsetarray.jl index d8e2518e2f139..4f8e5224db01b 100644 --- a/test/offsetarray.jl +++ b/test/offsetarray.jl @@ -609,3 +609,18 @@ end @test_throws DimensionMismatch maximum!(fill(0, -4:-4, 7:7, -6:-5, 1:1), B) @test_throws DimensionMismatch minimum!(fill(0, -4:-4, 7:7, -6:-5, 1:1), B) end + +@testset "first/last n elements of vector" begin + v0 = rand(6) + v = OffsetArray(v0, (-3,)) + @test_throws ArgumentError first(v, -2) + @test first(v, 2) == v[begin:begin+1] + @test first(v, 100) == v0 + @test first(v, 100) !== v + @test first(v, 1) == [v[begin]] + @test_throws ArgumentError last(v, -2) + @test last(v, 2) == v[end-1:end] + @test last(v, 100) == v0 + @test last(v, 100) !== v + @test last(v, 1) == [v[end]] +end