Merge pull request #129 from mcabbott/rf/broadcast-ambiguities

Fix #128 broadcast ambiguities

Project.toml

@@ -1,7 +1,7 @@
 name = "AxisKeys"
 uuid = "94b1ba4f-4ee9-5380-92f1-94cde586c3c5"
 license = "MIT"
-version = "0.2.8"
+version = "0.2.9"
 
 [deps]
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"

src/broadcast.jl

@@ -21,17 +21,28 @@ Base.BroadcastStyle(::Type{<:KeyedArray{T,N,AT}}) where {T,N,AT} =
 Base.BroadcastStyle(::KeyedStyle{A}, ::KeyedStyle{B}) where {A, B} = KeyedStyle(A(), B())
 Base.BroadcastStyle(::KeyedStyle{A}, b::B) where {A, B} = KeyedStyle(A(), b)
 Base.BroadcastStyle(a::A, ::KeyedStyle{B}) where {A, B} = KeyedStyle(a, B())
-Base.BroadcastStyle(::KeyedStyle{A}, b::DefaultArrayStyle) where {A} = KeyedStyle(A(), b)
-Base.BroadcastStyle(a::AbstractArrayStyle{M}, ::KeyedStyle{B}) where {B,M} = KeyedStyle(a, B())
 
 using NamedDims: NamedDimsStyle
 # this resolves in favour of KeyedArray(NamedDimsArray())
 Base.BroadcastStyle(a::NamedDimsStyle, ::KeyedStyle{B}) where {B} = KeyedStyle(a, B())
 Base.BroadcastStyle(::KeyedStyle{A}, b::NamedDimsStyle) where {A} = KeyedStyle(A(), b)
 
+# Resolve ambiguities
+# for all these cases, we define that we win to be the outer style regardless of order
+for B in (
+    :BroadcastStyle, :DefaultArrayStyle, :AbstractArrayStyle, :(Broadcast.Style{Tuple}),
+)
+    @eval function Base.BroadcastStyle(::KeyedStyle{A}, b::$B) where A
+        return KeyedStyle(A(), b)
+    end
+    @eval function Base.BroadcastStyle(b::$B, ::KeyedStyle{A}) where A
+        return KeyedStyle(b, A())
+    end
+end
+
 function unwrap_broadcasted(bc::Broadcasted{KeyedStyle{S}}) where {S}
     inner_args = map(unwrap_broadcasted, bc.args)
-    Broadcasted{S}(bc.f, inner_args)
+    Broadcasted{S}(bc.f, inner_args, axes(bc))
 end
 unwrap_broadcasted(x) = x
 unwrap_broadcasted(x::KeyedArray) = parent(x)

test/_basic.jl

@@ -259,6 +259,17 @@ end
         @test_throws Exception unify_keys((1:2,), ([3,4],))
 
     end
+    @testset "https://github.com/mcabbott/AxisKeys.jl/issues/128" begin
+        using LinearAlgebra
+
+        # Really simple test that broadcasting with linalg array types works.
+        v = rand(10)
+        σ = wrapdims(v; time=-4:5)
+        L = LowerTriangular(reshape(1.0:1.0:100, (10, 10)))
+        σ .* L
+        v .* L
+        @test σ .* L ≈ v .* L
+    end
 end
 @testset "bitarray" begin