Refactor similar methods

Author: MilesCranmer

src/arrays.jl

@@ -15,7 +15,7 @@ and so can be used in most places where a normal array would be used.
 # Constructors
 
 - `QuantityArray(value::AbstractArray, dimensions::AbstractDimensions)`: Create a `QuantityArray` with value `value` and dimensions `dimensions`.
-- `QuantityArray(value::AbstractArray, quantity::Quantity)`: Create a `QuantityArray` with value `value` and dimensions inferred 
+- `QuantityArray(value::AbstractArray, quantity::Quantity)`: Create a `QuantityArray` with value `value` and dimensions inferred
    with `dimension(quantity)`. This is so that you can easily create an array with the units module, like so:
    ```julia
    julia> A = QuantityArray(randn(32), 1u"m")
@@ -121,22 +121,19 @@ unsafe_setindex!(A, v, i...) = setindex!(ustrip(A), ustrip(v), i...)
 
 Base.IndexStyle(::Type{Q}) where {Q<:QuantityArray} = IndexStyle(array_type(Q))
 
-const IntOrOneTo = Union{Integer, Base.OneTo}
 
-# Unfortunately this mess of `similar` is required to avoid ambiguous methods.
-# c.f. base/abstractarray.jl
 Base.similar(A::QuantityArray) = QuantityArray(similar(ustrip(A)), dimension(A), quantity_type(A))
 Base.similar(A::QuantityArray, ::Type{S}) where {S} = QuantityArray(similar(ustrip(A), S), dimension(A), quantity_type(A))
-Base.similar(A::QuantityArray, dims::Dims{N}) where {N} = QuantityArray(similar(ustrip(A), dims), dimension(A), quantity_type(A))
-Base.similar(A::QuantityArray, dims::Tuple{IntOrOneTo, Vararg{IntOrOneTo}}) = QuantityArray(similar(ustrip(A), dims), dimension(A), quantity_type(A))
-Base.similar(A::QuantityArray, dims::Tuple{Integer, Vararg{Integer}}) = QuantityArray(similar(ustrip(A), dims), dimension(A), quantity_type(A))
-Base.similar(A::QuantityArray, ::Type{S}, dims::Dims{N}) where {S,N} = QuantityArray(similar(ustrip(A), S, dims), dimension(A), quantity_type(A))
-Base.similar(A::QuantityArray, ::Type{S}, dims::Tuple{IntOrOneTo, Vararg{IntOrOneTo}}) where {S} = QuantityArray(similar(ustrip(A), S, dims), dimension(A), quantity_type(A))
-Base.similar(A::QuantityArray, ::Type{S}, dims::Tuple{Integer, Vararg{Integer}}) where {S} = QuantityArray(similar(ustrip(A), S, dims), dimension(A), quantity_type(A))
-
-function Base.BroadcastStyle(::Type{QA}) where {QA<:QuantityArray}
-    return Broadcast.ArrayStyle{QA}()
+
+# Unfortunately this mess of `similar` is required to avoid ambiguous methods.
+# c.f. base/abstractarray.jl
+for dim_type in (:(Dims), :(Tuple{Union{Integer,Base.OneTo},Vararg{Union{Integer,Base.OneTo}}}), :(Tuple{Integer, Vararg{Integer}}))
+    @eval Base.similar(A::QuantityArray, dims::$dim_type) = QuantityArray(similar(ustrip(A), dims), dimension(A), quantity_type(A))
+    @eval Base.similar(A::QuantityArray, ::Type{S}, dims::$dim_type) where {S} = QuantityArray(similar(ustrip(A), S, dims), dimension(A), quantity_type(A))
 end
+
+Base.BroadcastStyle(::Type{QA}) where {QA<:QuantityArray} = Broadcast.ArrayStyle{QA}()
+
 function Base.similar(bc::Broadcast.Broadcasted{Broadcast.ArrayStyle{QA}}, ::Type{ElType}) where {QA<:QuantityArray,ElType<:AbstractQuantity}
     T = value_type(ElType)
     output_array = similar(bc, T)

test/unittests.jl

@@ -742,6 +742,14 @@ end
         a = [randn() * u"km/s" for i=1:32]
         @test typeof(y .* a) <: QuantityArray
         @test typeof(a .* y) <: QuantityArray
+
+        b = Quantity(randn(Float32, 32)) * u"km/s"
+        @test typeof(b) <: Quantity
+        @test typeof(b .* b) <: Vector{<:Quantity}
+        @test typeof(a .* b) <: Vector{<:Quantity}
+        @test typeof(b .* a) <: Vector{<:Quantity}
+        @test typeof(y .* b) <: QuantityArray{Float64}
+        @test typeof(b .* y) <: QuantityArray{Float64}
     end
 
     @testset "Broadcast scalars" begin
@@ -773,4 +781,16 @@ end
         msg2 = String(take!(io))
         @test msg2 == msg
     end
+
+    @testset "Extra test coverage" begin
+        @test_throws ErrorException DynamicQuantities.materialize_first(())
+        VERSION >= v"1.8" &&
+            @test_throws "Unexpected broadcast" DynamicQuantities.materialize_first(())
+
+        # Not sure how to test this otherwise, but method is supposed to be
+        # required for the broadcasting interface
+        x = [1u"km/s"]
+        ref = Base.RefValue(x)
+        @test DynamicQuantities.materialize_first(ref) === x[1]
+    end
 end