diff --git a/src/stack/stack.jl b/src/stack/stack.jl index 15f759511..b45179823 100644 --- a/src/stack/stack.jl +++ b/src/stack/stack.jl @@ -55,13 +55,24 @@ Base.similar(s::AbstractDimStack, args...) = map(A -> similar(A, args...), s) Base.eltype(s::AbstractDimStack, args...) = map(eltype, s) Base.iterate(s::AbstractDimStack, args...) = iterate(layers(s), args...) Base.read(s::AbstractDimStack) = map(read, s) +# `merge` for AbstractDimStack and NamedTuple. +# One of the first three arguments must be an AbstractDimStack for dispatch to work. Base.merge(s::AbstractDimStack) = s -function Base.merge(s1::AbstractDimStack, s2::AbstractDimStack, stacks::AbstractDimStack...) - rebuild_from_arrays(s1, merge(map(layers, (s1, s2, stacks...))...); refdims=()) +function Base.merge(x1::AbstractDimStack, x2::Union{AbstractDimStack,NamedTuple}, xs::Union{AbstractDimStack,NamedTuple}...) + rebuild_from_arrays(x1, merge(map(layers, (x1, x2, xs...))...); refdims=()) end function Base.merge(s::AbstractDimStack, pairs) rebuild_from_arrays(s, merge(layers(s), pairs); refdims=()) end +function Base.merge(x1::NamedTuple, x2::AbstractDimStack, xs::Union{AbstractDimStack,NamedTuple}...) + merge(map(layers, (x1, x2, xs...))...) +end +function Base.merge(x1::NamedTuple, x2::NamedTuple, x3::AbstractDimStack, xs::Union{AbstractDimStack,NamedTuple}...) + merge(map(layers, (x1, x2, x3, xs...))...) +end +function Base.setindex(s::AbstractDimStack, val::AbstractDimArray, key) + rebuild_from_arrays(s, Base.setindex(layers(s), val, key)) +end Base.NamedTuple(s::AbstractDimStack) = layers(s) Extents.extent(A::AbstractDimStack, args...) = Extents.extent(dims(A), args...) @@ -113,6 +124,7 @@ function rebuild_from_arrays( rebuild(s; data, dims, refdims, layerdims, metadata, layermetadata) end +layers(nt::NamedTuple) = nt function layers(s::AbstractDimStack{<:NamedTuple{Keys}}) where Keys NamedTuple{Keys}(map(K -> s[K], Keys)) end diff --git a/test/stack.jl b/test/stack.jl index 838f44f22..b433814de 100644 --- a/test/stack.jl +++ b/test/stack.jl @@ -3,7 +3,7 @@ using DimensionalData, Test, LinearAlgebra, Statistics using DimensionalData: data using DimensionalData: Sampled, Categorical, AutoLookup, NoLookup, Transformed, Regular, Irregular, Points, Intervals, Start, Center, End, - Metadata, NoMetadata, ForwardOrdered, ReverseOrdered, Unordered + Metadata, NoMetadata, ForwardOrdered, ReverseOrdered, Unordered, layers A = [1.0 2.0 3.0; 4.0 5.0 6.0] @@ -74,12 +74,24 @@ end end @testset "merge" begin + @test layers(s) === (; s...) + @test layers(s) === (; one=da1, s[(:two,)]..., (three=da3,)...) + @test layers(s) === (; (one=da1,)..., two=da2, s[(:three,)]...) + @test s === merge(s[(:one,:two)], (three=da3,)) + @test s === merge(s[(:one,)], (two=da2, three=da3)) @test merge(mixed) === mixed @test keys(merge(mixed, s)) == (:one, :two, :extradim, :three) @test keys(merge(s, mixed)) == (:one, :two, :three, :extradim) @test keys(merge(s, (:new=>da4,))) == (:one, :two, :three, :new) end +@testset "setindex" begin + s12 = DimStack((da1, da2)) + @test s === Base.setindex(s12, da3, :three) + s423 = DimStack((one=da4, two=da2, three=da3)) + @test s === Base.setindex(s423, da1, :one) +end + @testset "copy and friends" begin sc = copy(s) @test sc == s