instead, implement Random.Sampler and Rand.rand(rng, ::Sampler{...})

Author: ablaom

src/methods.jl

@@ -260,27 +260,26 @@ function _rand(rng, p_cumulative, R)
     return index
 end
 
-function Base.rand(rng::AbstractRNG,
-                   d::UnivariateFinite{<:Any,<:Any,R}) where R
-    p_cumulative = _cumulative(d)
-    return Dist.support(d)[_rand(rng, p_cumulative, R)]
-end
-Base.rand(d::UnivariateFinite) = rand(Random.default_rng(), d)
-
-function Base.rand(rng::AbstractRNG,
-                   d::UnivariateFinite{<:Any,<:Any,R},
-                   dim1::Integer, moredims::Integer...) where R # ref type
-    p_cumulative = _cumulative(d)
-    A = Array{R}(undef, dim1, moredims...)
-    for i in eachindex(A)
-        @inbounds A[i] = _rand(rng, p_cumulative, R)
-    end
-    support = Dist.support(d)
-    return broadcast(i -> support[i], A)
+Random.eltype(::Type{<:UnivariateFinite{<:Any,V}}) where V = V
+
+# The Sampler hook into Random's API is discussed in the Julia documentation, in the
+# Standard Library section on Random.
+function Random.Sampler(
+    ::AbstractRNG,
+    d::UnivariateFinite,
+    ::Random.Repetition,
+    )
+    data = (_cumulative(d), Dist.support(d))
+    Random.SamplerSimple(d, data)
 end
 
-Base.rand(d::UnivariateFinite, dim1::Integer, moredims::Integer...) =
-    rand(Random.default_rng(), d, dim1, moredims...)
+function Base.rand(
+    rng::AbstractRNG,
+    sampler::Random.SamplerSimple{<:UnivariateFinite{<:Any,<:Any,R}},
+    ) where R
+    p_cumulative, support = sampler.data
+    return support[_rand(rng, p_cumulative, R)]
+end
 
 function Dist.fit(d::Type{<:UnivariateFinite},
                            v::AbstractVector{C}) where C

test/methods.jl

@@ -300,24 +300,26 @@ end
     @test displays_okay([5 + 3im, 4 - 7im])
 end
 
-if VERSION >= v"1.7"
-    @testset "rand signatures" begin
-        d = UnivariateFinite(
-            ["maybe", "no", "yes"],
-            [0.5, 0.4, 0.1];
-            pool=missing,
-        )
-
-        Random.seed!(123)
-        samples = [rand(default_rng(), d) for i in 1:30]
-        Random.seed!(123)
-        @test [rand(d) for i in 1:30] == samples
-
-        Random.seed!(123)
-        samples = rand(Random.default_rng(), d, 3, 5)
-        Random.seed!(123)
-        @test samples == rand(d, 3, 5)
-    end
+@testset "rand signatures" begin
+    d = UnivariateFinite(
+        ["maybe", "no", "yes"],
+        [0.5, 0.4, 0.1];
+        pool=missing,
+    )
+
+    # smoke test:
+    sampler = Random.Sampler(default_rng(), d, Val(1))
+    rand(default_rng(), sampler)
+
+    Random.seed!(123)
+    samples = [rand(default_rng(), d) for i in 1:30]
+    Random.seed!(123)
+    @test [rand(d) for i in 1:30] == samples
+
+    Random.seed!(123)
+    samples = rand(Random.default_rng(), d, 3, 5)
+    Random.seed!(123)
+    @test samples == rand(d, 3, 5)
 end
 
 end # module