Prevent overflow in mean(::AbstractRange) and relax type constraint. (#150)

Co-authored-by: Chunji Wang <[email protected]>
Co-authored-by: Vincent Yu <[email protected]>

Author: 3 people

src/Statistics.jl

@@ -196,9 +196,9 @@ function _mean(f, A::AbstractArray, dims::Dims=:) where Dims
     end
 end
 
-function mean(r::AbstractRange{<:Real})
-    isempty(r) && return oftype((first(r) + last(r)) / 2, NaN)
-    (first(r) + last(r)) / 2
+function mean(r::AbstractRange{T}) where T
+    isempty(r) && return zero(T)/0
+    return first(r)/2 + last(r)/2
 end
 
 ##### variances #####

test/runtests.jl

@@ -194,8 +194,16 @@ end
             @test f(2:0.1:n) ≈ f([2:0.1:n;])
         end
     end
-    @test mean(2:1) === NaN
-    @test mean(big(2):1) isa BigFloat
+    @test mean(2:0.1:4) === 3.0  # N.B. mean([2:0.1:4;]) != 3
+    @test mean(LinRange(2im, 4im, 21)) === 3.0im
+    @test mean(2:1//10:4) === 3//1
+    @test isnan(@inferred(mean(2:1))::Float64)
+    @test isnan(@inferred(mean(big(2):1))::BigFloat)
+    z = @inferred(mean(LinRange(2im, 1im, 0)))::ComplexF64
+    @test isnan(real(z)) & isnan(imag(z))
+    @test_throws DivideError mean(2//1:1)
+    @test mean(typemax(Int):typemax(Int)) === float(typemax(Int))
+    @test mean(prevfloat(Inf):prevfloat(Inf)) === prevfloat(Inf)
 end
 
 @testset "var & std" begin