add setindex and improve merge for stacks

rafaqz · Sep 26, 2022 · 55c0331 · 55c0331
1 parent 455d958
commit 55c0331
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 3 deletions.
diff --git 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
@@ -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