Merge pull request #158 from SymbolicML/MilesCranmer/issue157

feat: add explicit handling of ranges

Author: MilesCranmer

src/utils.jl

@@ -149,6 +149,34 @@ Base.getindex(d::AbstractDimensions, k::Symbol) = getfield(d, k)
     return dimension_names(T1) == dimension_names(T2)
 end
 
+# Multiplying ranges with units
+Base.TwicePrecision{T}(x::T) where {T<:AbstractQuantity} = Base.TwicePrecision{typeof(x)}(x, zero(x))
+# TODO: Note that to get RealQuantity working, we have to overload many other functions,
+#       which is why we skip it.
+
+# Avoid https://github.com/JuliaLang/julia/issues/56610
+for T1 in (AbstractQuantity{<:Real}, Real),
+    T2 in (AbstractQuantity{<:Real}, Real),
+    T3 in (AbstractQuantity{<:Real}, Real)
+
+    T1 === T2 === T3 === Real && continue
+
+    @eval function Base.:(:)(start::$T1, step::$T2, stop::$T3)
+        dimension(start) == dimension(step) || throw(DimensionError(start, step))
+        dimension(start) == dimension(stop) || throw(DimensionError(start, stop))
+        return range(start, stop, length=length(ustrip(start):ustrip(step):ustrip(stop)))
+    end
+
+    if T3 === Real && !(T1 === T2 === Real)
+        @eval function Base.:(:)(start::$T1, stop::$T2)
+            if !iszero(dimension(start)) || !iszero(dimension(stop))
+                error("When creating a range over dimensionful quantities, you must specify a step.")
+            end
+            return start:1:stop
+        end
+    end
+end
+
 # Compatibility with `.*`
 Base.size(q::UnionAbstractQuantity) = size(ustrip(q))
 Base.length(q::UnionAbstractQuantity) = length(ustrip(q))

test/unittests.jl

@@ -198,12 +198,6 @@ Base.round(::Type{T}, x::SimpleRatio) where {T} = round(T, x.num // x.den)
 
 end
 
-@testset "Ranges" begin
-    x = [xi for xi in 0.0u"km/s":0.1u"km/s":1.0u"km/s"]
-    @test x[2] == 0.1u"km/s"
-    @test x[end] == 1.0u"km/s"
-end
-
 @testset "Complex numbers" begin
     x = (0.5 + 0.6im) * u"km/s"
     @test string(x) == "(500.0 + 600.0im) m s⁻¹"
@@ -358,6 +352,109 @@ end
     @test ustrip(x') == ustrip(x)'
 end
 
+@testset "Ranges" begin
+    @testset "Ranges from units" begin
+        x = [xi for xi in 0.0u"km/s":0.1u"km/s":1.0u"km/s"]
+        @test x[2] == 0.1u"km/s"
+        @test x[end] == 1.0u"km/s"
+
+        # https://github.com/JuliaLang/julia/issues/56610
+        c = collect(1u"inch":0.25u"inch":4u"inch")
+        @test c[1] == 1u"inch"
+        @test c[end] <= 4u"inch"
+
+        # Test dimensionless quantities
+        x = collect(1u"1":2u"1":5u"1")
+        @test x == [1, 3, 5] .* u"1"
+        @test eltype(x) <: Quantity
+
+        # Test error for missing step
+        @test_throws "must specify a step" 1u"km":2u"km"
+        @test_throws "must specify a step" 1us"km":2us"km"
+
+        # However, for backwards compatibility, dimensionless ranges are allowed:
+        x = collect(1u"1":5u"1")
+        @test x == [1, 2, 3, 4, 5]
+        @test eltype(x) <: Quantity{Float64}
+
+        # Test errors for incompatible units
+        @test_throws DimensionError 1u"km":1u"s":5u"km"
+        @test_throws DimensionError 1u"km":1u"m":5u"s"
+        @test_throws DimensionError 1u"km":1u"km/s":5u"km"
+
+        # Same for symbolic units!
+        @test_throws DimensionError 1us"km":1us"m":5us"inch"
+        @test length(1u"inch":1u"m":5u"km") == 5000
+
+        # Test with symbolic units
+        x = collect(1us"inch":0.25us"inch":4us"inch")
+        @test x[1] == 1us"inch"
+        @test x[2] == 1.25us"inch"
+        @test x[end] == 4us"inch"
+    end
+
+    @testset "Multiplying ranges with units" begin
+        # Test multiplying ranges with units
+        x = (1:0.25:4)u"inch"
+        @test x isa StepRangeLen
+        @test first(x) == 1u"inch"
+        @test x[2] == 1.25u"inch"
+        @test last(x) == 4u"inch"
+        @test length(x) == 13
+
+        # Integer range (but real-valued unit)
+        x = (1:4)u"inch"
+        @test x isa StepRangeLen
+        @test first(x) == 1u"inch"
+        @test x[2] == 2u"inch"
+        @test last(x) == 4u"inch"
+        @test length(x) == 4
+        @test collect(x)[3] == 3u"inch"
+
+        # Test with floating point range
+        x = (1.0:0.5:3.0)u"m"
+        @test x isa StepRangeLen
+        @test first(x) == 1.0u"m"
+        @test x[2] == 1.5u"m"
+        @test last(x) == 3.0u"m"
+        @test length(x) == 5
+        @test collect(x)[3] == 2.0u"m"
+
+        x = (0:0.1:1)u"m"
+        @test length(x) == 11
+        @test collect(x)[3] == 0.2u"m"
+
+        # Test with symbolic units
+        x = (1:0.25:4)us"inch"
+        @test x isa StepRangeLen{<:Quantity{Float64,<:SymbolicDimensions}}
+        @test first(x) == us"inch"
+        @test x[2] == 1.25us"inch"
+        @test last(x) == 4us"inch"
+        @test length(x) == 13
+
+        # Test that symbolic units preserve their symbolic nature
+        x = (0:0.1:1)us"km/h"
+        @test x isa AbstractRange
+        @test first(x) == 0us"km/h"
+        @test x[2] == 0.1us"km/h"
+        @test last(x) == 1us"km/h"
+        @test length(x) == 11
+
+        # Similarly, integers should stay integers:
+        x = (1:4)us"inch"
+        @test_skip x isa StepRangeLen{<:Quantity{Int64,<:SymbolicDimensions}}
+        @test first(x) == us"inch"
+        @test x[2] == 2us"inch"
+        @test last(x) == 4us"inch"
+        @test length(x) == 4
+
+        # With RealQuantity:
+        @test_skip (1.0:4.0) * RealQuantity(u"inch") isa StepRangeLen{<:RealQuantity{Float64,<:SymbolicDimensions}}
+        # TODO: This is not available as TwicePrecision interacts with Real in a way
+        #       that demands many other functions to be defined.
+    end
+end
+
 @testset "Alternate dimension construction" begin
     z = Quantity(-52, length=1) * Dimensions(mass=2)
     z2 = Dimensions(mass=2) * Quantity(-52, length=1)