From 16266fa623319c073bdc31be01bf49d4a59093dd Mon Sep 17 00:00:00 2001 From: Dilum Aluthge Date: Tue, 1 Mar 2022 13:19:43 -0500 Subject: [PATCH] Add the `only(f, x)` method, which allows users to provide a default value or customize the error message --- base/iterators.jl | 47 +++++++++++++++++++++--- test/iterators.jl | 91 ++++++++++++++++++++++++++++++++++++++++++++++- 2 files changed, 132 insertions(+), 6 deletions(-) diff --git a/base/iterators.jl b/base/iterators.jl index 1b96a24a9c16f..cb44f199019f5 100644 --- a/base/iterators.jl +++ b/base/iterators.jl @@ -1373,14 +1373,20 @@ length(s::Stateful) = length(s.itr) - s.taken """ only(x) - Return the one and only element of collection `x`, or throw an [`ArgumentError`](@ref) if the collection has zero or multiple elements. + only(f, x) +Return the one and only element of collection `x`, or return `f()` if the collection has zero +or multiple elements. `f()` will not be called if the collection has exactly one element. + See also [`first`](@ref), [`last`](@ref). !!! compat "Julia 1.4" - This method requires at least Julia 1.4. + The `only(x)` method requires at least Julia 1.4. + +!!! compat "Julia 1.9" + The `only(f, x) method requires at least Julia 1.9. # Examples ```jldoctest @@ -1399,6 +1405,16 @@ julia> only(('a', 'b')) ERROR: ArgumentError: Tuple contains 2 elements, must contain exactly 1 element Stacktrace: [...] + +julia> only(('a', 'b')) do + 'c' + end +'c': ASCII/Unicode U+0063 (category Ll: Letter, lowercase) + +julia> only(('a', 'b')) do + throw(ErrorException("My custom error message")) + end +ERROR: My custom error message ``` """ @propagate_inbounds function only(x) @@ -1412,20 +1428,41 @@ Stacktrace: end return ret end +function only(f::F, x) where {F} + i = iterate(x) + if i === nothing + return f() + end + (ret, state) = i::NTuple{2,Any} + if iterate(x, state) !== nothing + return f() + end + return ret +end -# Collections of known size +# Collections of known size that have exactly one element only(x::Ref) = x[] only(x::Number) = x only(x::Char) = x only(x::Tuple{Any}) = x[1] +only(a::AbstractArray{<:Any, 0}) = @inbounds return a[] +only(x::NamedTuple{<:Any, <:Tuple{Any}}) = first(x) +only(f::F, x::Ref) where {F} = only(x) +only(f::F, x::Number) where {F} = only(x) +only(f::F, x::Char) where {F} = only(x) +only(f::F, x::Tuple{Any}) where {F} = only(x) +only(f::F, a::AbstractArray{<:Any, 0}) where {F} = only(a) +only(f::F, x::NamedTuple{<:Any, <:Tuple{Any}}) where {F} = only(x) + +# Collections of known size that either have zero elements or more than one element only(x::Tuple) = throw( ArgumentError("Tuple contains $(length(x)) elements, must contain exactly 1 element") ) -only(a::AbstractArray{<:Any, 0}) = @inbounds return a[] -only(x::NamedTuple{<:Any, <:Tuple{Any}}) = first(x) only(x::NamedTuple) = throw( ArgumentError("NamedTuple contains $(length(x)) elements, must contain exactly 1 element") ) +only(f::F, x::Tuple) where {F} = f() +only(f::F, x::NamedTuple) where {F} = f() Base.intersect(a::ProductIterator, b::ProductIterator) = ProductIterator(intersect.(a.iterators, b.iterators)) diff --git a/test/iterators.jl b/test/iterators.jl index 70ce6866f4be3..e7ed2088edc0c 100644 --- a/test/iterators.jl +++ b/test/iterators.jl @@ -832,7 +832,7 @@ end @test length(collect(d)) == 0 end -@testset "only" begin +@testset "only(x)" begin @test only([3]) === 3 @test_throws ArgumentError only([]) @test_throws ArgumentError only([3, 2]) @@ -865,6 +865,95 @@ end @test_throws ArgumentError only(1 for ii in 1:10 if ii > 200) end +@testset "only(f, x)" begin + my_exception = ErrorException("Hello world") + my_throw = () -> throw(my_exception) + @test 3 === only([3]) do + my_throw() + end + @test_throws my_exception only([]) do + my_throw() + end + @test_throws ErrorException("Hello world") only([3, 2]) do + my_throw() + end + + f = () -> only((3,)) do + my_throw() + end + @test 3 === @inferred(f()) + @test_throws my_exception only(()) do + my_throw() + end + @test_throws my_exception only((3, 2)) do + my_throw() + end + + @test (1=>3) === only(Dict(1=>3)) do + my_throw() + end + @test_throws my_exception only(Dict{Int,Int}()) do + my_throw() + end + @test_throws my_exception only(Dict(1=>3, 2=>2)) do + my_throw() + end + + @test 3 === only(Set([3])) do + my_throw() + end + @test_throws my_exception only(Set(Int[])) do + my_throw() + end + @test_throws my_exception only(Set([3,2])) do + my_throw() + end + + f = () -> only((;a=1)) do + my_throw() + end + @test 1 === @inferred(f()) + @test_throws my_exception only(NamedTuple()) do + my_throw() + end + @test_throws my_exception only((a=3, b=2.0)) do + my_throw() + end + + f = () -> only(1) do + my_throw() + end + @test 1 === @inferred(f()) + f = () -> only('a') do + my_throw() + end + @test 'a' === @inferred(f()) + f = () -> only(Ref([1, 2])) do + my_throw() + + end + @test [1, 2] == @inferred(f()) + @test_throws my_exception only(Pair(10, 20)) do + my_throw() + end + + @test 1 === only(1 for ii in 1:1) do + my_throw() + end + @test 1 === only(1 for ii in 1:10 if ii < 2) do + my_throw() + end + @test_throws my_exception only(1 for ii in 1:10) do + my_throw() + end + @test_throws my_exception only(1 for ii in 1:10 if ii > 2) do + my_throw() + end + @test_throws my_exception only(1 for ii in 1:10 if ii > 200) do + my_throw() + end +end + @testset "flatten empty tuple" begin @test isempty(collect(Iterators.flatten(()))) end