dont use nothing for layerdims

rafaqz · Apr 1, 2023 · 91c5b2d · 91c5b2d
1 parent ba98b75
commit 91c5b2d
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 8 deletions.
diff --git a/src/stack/stack.jl b/src/stack/stack.jl
@@ -33,9 +33,7 @@ refdims(s::AbstractDimStack) = getfield(s, :refdims)
 metadata(s::AbstractDimStack) = getfield(s, :metadata)
 
 layerdims(s::AbstractDimStack) = getfield(s, :layerdims)
-function layerdims(s::AbstractDimStack, key::Symbol)
-    isnothing(layerdims(s)) ? dims(s) : dims(s, layerdims(s)[key])
-end
+layerdims(s::AbstractDimStack, key::Symbol) = dims(s, layerdims(s)[key])
 layermetadata(s::AbstractDimStack) = getfield(s, :layermetadata)
 layermetadata(s::AbstractDimStack, key::Symbol) = layermetadata(s)[key]
 
@@ -255,8 +253,7 @@ function DimStack(@nospecialize(das::AbstractArray{<:AbstractDimArray});
     dims = DD.combinedims(collect(das))
     as = map(parent, das)
     data = NamedTuple{keys_tuple}(as)
-    same_dims_layers = all(map(a -> ndims(a) == length(dims), as))
-    layerdims = same_dims_layers ? nothing : NamedTuple{keys_tuple}(map(basedims, das))
+    layerdims = NamedTuple{keys_tuple}(map(basedims, das))
     layermetadata = NamedTuple{keys_tuple}(map(DD.metadata, das))
 
     DimStack(data, dims, refdims, layerdims, metadata, layermetadata)
@@ -269,10 +266,11 @@ function DimStack(das::NamedTuple{<:Any,<:Tuple{Vararg{AbstractDimArray}}};
 end
 # Same sized arrays
 function DimStack(data::NamedTuple, dims::Tuple;
-    refdims=(), metadata=NoMetadata(), layermetadata=map(_ -> NoMetadata(), data)
+    refdims=(), metadata=NoMetadata(), 
+    layermetadata=map(_ -> NoMetadata(), data),
+    layerdims = map(_ -> basedims(dims), data),
 )
     all(map(d -> axes(d) == axes(first(data)), data)) || _stack_size_mismatch()
-    layerdims = nothing
     DimStack(data, format(dims, first(data)), refdims, layerdims, metadata, layermetadata)
 end
 

diff --git a/test/stack.jl b/test/stack.jl
@@ -96,7 +96,7 @@ end
     @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 typeof(s) === typeof(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)