Add the only(f, x) method, which allows users to provide a default value or customize the error message

Author: DilumAluthge

base/iterators.jl

@@ -1413,19 +1413,68 @@ Stacktrace:
     return ret
 end
 
-# Collections of known size
+"""
+    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.
+
+!!! compat "Julia 1.9"
+    This method requires at least Julia 1.9.
+
+# Examples
+```jldoctest
+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
+```
+"""
+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 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))

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,137 @@ 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 1 === only([]) do
+        1
+    end
+    @test_throws my_exception only([3, 2]) do
+        my_throw()
+    end
+    @test 1 === only([3, 2]) do
+        1
+    end
+
+    f = () -> only((3,)) do
+        my_throw()
+    end
+    @test 3 === @inferred(f())
+    @test_throws my_exception only(()) do
+        my_throw()
+    end
+    @test 1 === only(()) do
+        1
+    end
+    @test_throws my_exception only((3, 2)) do
+        my_throw()
+    end
+    @test 1 === only((3, 2)) do
+        1
+    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 1 === only(Dict{Int,Int}()) do
+        1
+    end
+    @test_throws my_exception only(Dict(1=>3, 2=>2)) do
+        my_throw()
+    end
+    @test 1 === only(Dict(1=>3, 2=>2)) do
+        1
+    end
+
+    @test 3 === only(Set([3])) do
+        my_throw()
+    end
+    @test_throws my_exception only(Set(Int[])) do
+        my_throw()
+    end
+    @test 1 === only(Set(Int[])) do
+        1
+    end
+    @test_throws my_exception only(Set([3,2])) do
+        my_throw()
+    end
+    @test 1 === only(Set([3,2])) do
+        1
+    end
+
+    f = () -> only((;a=1)) do
+        my_throw()
+    end
+    @test 1 === @inferred(f())
+    @test_throws my_exception only(NamedTuple()) do
+        my_throw()
+    end
+    @test 1 === only(NamedTuple()) do
+        1
+    end
+    @test_throws my_exception only((a=3, b=2.0)) do
+        my_throw()
+    end
+    @test 1 === only((a=3, b=2.0)) do
+        1
+    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 30 === only(Pair(10, 20)) do
+        30
+    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 0 === only(1 for ii in 1:10) do
+        0
+    end
+    @test_throws my_exception only(1 for ii in 1:10 if ii > 2) do
+        my_throw()
+    end
+    @test 0 === only(1 for ii in 1:10 if ii > 2) do
+        0
+    end
+    @test_throws my_exception only(1 for ii in 1:10 if ii > 200) do
+        my_throw()
+    end
+    @test 0 === only(1 for ii in 1:10 if ii > 200) do
+        0
+    end
+end
+
 @testset "flatten empty tuple" begin
     @test isempty(collect(Iterators.flatten(())))
 end