fix: handle broadcast allocation for range-backed QuantityArray (#218)

* fix: handle broadcast allocation for range-backed QuantityArray

Co-authored-by: Miles Cranmer <miles.cranmer@gmail.com>

* refactor: use @eval loop for quantity array similar

Co-authored-by: Miles Cranmer <miles.cranmer@gmail.com>

---------

Co-authored-by: Miles Cranmer <miles.cranmer@gmail.com>

Author: MilesCranmerBot

src/arrays.jl

@@ -315,8 +315,13 @@ function Base.similar(bc::Broadcast.Broadcasted{Broadcast.ArrayStyle{QA}}, ::Typ
     first_output::ElType = materialize_first(bc)
     return QuantityArray(output_array, dimension(first_output)::dim_type(ElType), ElType)
 end
-function Base.similar(bc::Broadcast.Broadcasted{Broadcast.ArrayStyle{QuantityArray{T,N,D,Q,V}}}, ::Type{ElType}) where {T,N,D,Q,V<:Array{T,N},ElType}
-    return similar(Array{ElType}, axes(bc))
+for VSupertype in (:(Array{T,N}), :(AbstractRange{T}))
+    @eval function Base.similar(
+        bc::Broadcast.Broadcasted{Broadcast.ArrayStyle{QuantityArray{T,N,D,Q,V}}},
+        ::Type{ElType},
+    ) where {T,N,D,Q,V<:$VSupertype,ElType}
+        return similar(Array{ElType}, axes(bc))
+    end
 end
 function Base.similar(bc::Broadcast.Broadcasted{Broadcast.ArrayStyle{QuantityArray{T,N,D,Q,V}}}, ::Type{ElType}) where {T,N,D,Q,V,ElType}
     # To deal with things like StaticArrays, we need to rely on

test/unittests.jl

@@ -482,6 +482,19 @@ end
         @test RealQuantity(us"inch") * (1.0:4.0) isa QuantityArray
     end
 
+    @testset "Broadcasting over QuantityArray with StepRangeLen backing" begin
+        struct WrappedQuantity{T}
+            value::T
+        end
+
+        t = QuantityArray(range(1.0, step=0.1, length=10), u"s")
+        wrapped = WrappedQuantity.(t)
+
+        @test wrapped isa Vector{WrappedQuantity{Quantity{Float64,typeof(dimension(u"s"))}}}
+        @test first(wrapped).value == first(t)
+        @test last(wrapped).value == last(t)
+    end
+
     @testset "Array of quantities multiplication behavior" begin
         # Test that array of quantities * quantity does NOT create nested QuantityArray
         # but instead broadcasts properly