Skip to content

Commit

Permalink
add get method to AbstractDimStack (#889)
Browse files Browse the repository at this point in the history
  • Loading branch information
rafaqz authored Dec 29, 2024
1 parent 688e994 commit 610f2f9
Show file tree
Hide file tree
Showing 2 changed files with 8 additions and 1 deletion.
4 changes: 4 additions & 0 deletions src/stack/stack.jl
Original file line number Diff line number Diff line change
Expand Up @@ -189,6 +189,10 @@ Base.NamedTuple(s::AbstractDimStack) = NamedTuple(layers(s))
Base.collect(st::AbstractDimStack) = parent([st[D] for D in DimIndices(st)])
Base.Array(st::AbstractDimStack) = collect(st)
Base.vec(st::AbstractDimStack) = vec(collect(st))
Base.get(st::AbstractDimStack, k::Symbol, default) =
haskey(st, k) ? st[k] : default
Base.get(f::Base.Callable, st::AbstractDimStack, k::Symbol) =
haskey(st, k) ? st[k] : f()
@propagate_inbounds Base.iterate(st::AbstractDimStack) = iterate(st, 1)
@propagate_inbounds Base.iterate(st::AbstractDimStack, i) =
i > length(st) ? nothing : (st[DimIndices(st)[i]], i + 1)
Expand Down
5 changes: 4 additions & 1 deletion test/stack.jl
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,6 @@ da4 = DimArray(cat(4A, 5A, 6A, 7A; dims=3), (x, y, z); name=:extradim)

s = DimStack((da1, da2, da3))
mixed = DimStack(da1, da2, da4)
typeof(s)

@testset "constructors" begin
@test DimStack((one=A, two=2A, three=3A), dimz) ==
Expand Down Expand Up @@ -67,6 +66,10 @@ end
@test eltype(mixed) === @NamedTuple{one::Float64, two::Float32, extradim::Float64}
@test haskey(s, :one) == true
@test haskey(s, :zero) == false
@test get(mixed, :one, nothing) == mixed.one
@test get(mixed, :five, nothing) == nothing
@test get(() -> nothing, mixed, :two) == mixed.two
@test get(() -> nothing, mixed, :ten) == nothing
@test size(mixed) === (2, 3, 4) # size is as for Array
@test size(mixed, Y) === 3
@test size(mixed, 3) === 4
Expand Down

0 comments on commit 610f2f9

Please sign in to comment.