Add the only(f, x) method, which allows users to provide a default …

…value or customize the error message

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))

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