Skip to content

Commit 8d73d2a

Browse files
aplavinmcabbott
andauthored
fixes and updates for Julia 1.9 (#130)
* update uncolon * eachslice with 1.9 semantics * add stack * add TODO * fix tests on 1.9 * properly qualify stack * remove accidental piracy * Update src/functions.jl --------- Co-authored-by: Michael Abbott <32575566+mcabbott@users.noreply.github.com>
1 parent 7929518 commit 8d73d2a

File tree

7 files changed

+55
-19
lines changed

7 files changed

+55
-19
lines changed

src/functions.jl

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -94,7 +94,13 @@ function Base.permutedims(A::KeyedArray, perm)
9494
KeyedArray(data, new_keys)#, copy(A.meta))
9595
end
9696

97-
if VERSION >= v"1.1"
97+
@static if VERSION > v"1.9-DEV"
98+
function Base.eachslice(A::KeyedArray; dims)
99+
dims_ix = AxisKeys.dim(A, dims) |> Tuple
100+
data = @invoke eachslice(A::AbstractArray; dims=dims_ix)
101+
return KeyedArray(NamedDimsArray(data, map(d -> dimnames(A, d), dims_ix)), map(d -> axiskeys(A, d), dims_ix))
102+
end
103+
elseif VERSION >= v"1.1"
98104
# This copies the implementation from Base, except with numerical_dims:
99105
@inline function Base.eachslice(A::KeyedArray; dims)
100106
numerical_dims = NamedDims.dim(A, dims)
@@ -107,6 +113,19 @@ if VERSION >= v"1.1"
107113
end
108114
end
109115

116+
@static if VERSION > v"1.9-DEV"
117+
# TODO: this will ERROR if given dims, instead of falling back to Base
118+
# TODO: ideally it would dispatch on the element type, for e.g. a generator of KeyedArrays
119+
function Base.stack(A::KeyedArray; dims::Colon=:)
120+
data = @invoke stack(A::AbstractArray; dims)
121+
if !allequal(named_axiskeys(a) for a in A)
122+
throw(DimensionMismatch("stack expects uniform axiskeys for all arrays"))
123+
end
124+
akeys = (; named_axiskeys(first(A))..., named_axiskeys(A)...)
125+
KeyedArray(data; akeys...)
126+
end
127+
end
128+
110129
function Base.mapslices(f, A::KeyedArray; dims)
111130
numerical_dims = NamedDims.dim(A, dims)
112131
data = mapslices(f, parent(A); dims=numerical_dims)

src/selectors.jl

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -208,8 +208,14 @@ end
208208
@inline Base.to_indices(A::Union{KeyedArray,NdaKa}, ax, inds::Tuple{Function, Vararg}) =
209209
select_to_indices(A, ax, inds)
210210

211-
using Base: to_indices, tail, safe_tail, uncolon
211+
using Base: to_indices, tail, safe_tail
212+
213+
if VERSION > v"1.9-DEV"
214+
const uncolon = Base.uncolon
215+
else
216+
uncolon(inds) = Base.uncolon(inds, (:,))
217+
end
212218

213219
@inline Base.to_indices(A::Union{KeyedArray,NdaKa}, inds, I::Tuple{Colon, Vararg{Any}}) =
214-
(uncolon(inds, I), to_indices(A, safe_tail(inds), tail(I))...)
220+
(uncolon(inds), to_indices(A, safe_tail(inds), tail(I))...)
215221

src/stack.jl

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11

2-
using LazyStack
2+
import LazyStack
33

44
# for stack_iter
55
LazyStack.no_wraps(a::KeyedArray) = LazyStack.no_wraps(NamedDims.unname(parent(a)))
@@ -23,14 +23,14 @@ stack_keys(xs::Tuple{Vararg{KeyedArray}}) =
2323

2424
# array of arrays: first strip off outer containers...
2525
function LazyStack.stack(xs::KeyedArray{<:AbstractArray})
26-
KeyedArray(stack(parent(xs)), stack_keys(xs))
26+
KeyedArray(LazyStack.stack(parent(xs)), stack_keys(xs))
2727
end
2828
function LazyStack.stack(xs::KeyedArray{<:AbstractArray,N,<:NamedDimsArray{L}}) where {L,N}
29-
data = stack(parent(parent(xs)))
29+
data = LazyStack.stack(parent(parent(xs)))
3030
KeyedArray(LazyStack.ensure_named(data, LazyStack.getnames(xs)), stack_keys(xs))
3131
end
3232
function LazyStack.stack(xs::NamedDimsArray{L,<:AbstractArray,N,<:KeyedArray}) where {L,N}
33-
data = stack(parent(parent(xs)))
33+
data = LazyStack.stack(parent(parent(xs)))
3434
LazyStack.ensure_named(KeyedArray(data, stack_keys(xs)), LazyStack.getnames(xs))
3535
end
3636

test/_basic.jl

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -151,7 +151,7 @@ end
151151
N2 = NamedDimsArray(KeyedArray(data, axiskeys(N1)), dimnames(N1))
152152
N3 = KeyedArray(NamedDimsArray(data, dimnames(N1)), axiskeys(N1))
153153

154-
@testset "with $(typeof(N).name) outside" for N in [N2, N3]
154+
@testset "with $(typeof(N).name) outside" for N in (N2, N3)
155155
@test axiskeys(N) == (['a', 'b', 'c'], 10:10:40)
156156
@test axiskeys(N, :iter) == 10:10:40
157157
@test dimnames(N) == (:obs, :iter)

test/_functions.jl

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,17 @@ A3 = wrapdims(rand(Int8, 3,4,2), r='a':'c', c=2:5, p=[10.0, 20.0])
4646
if VERSION >= v"1.1"
4747
@test axiskeys(first(eachslice(M, dims=:r))) === (2:5,)
4848
end
49+
if VERSION >= v"1.9-DEV"
50+
arr = KeyedArray([
51+
KeyedArray([1, 2], a=[:x, :y]),
52+
KeyedArray([3, 4], a=[:x, :y]),
53+
KeyedArray([5, 6], a=[:x, :y]),
54+
], b=10:12)
55+
56+
sk = stack(arr)::KeyedArray
57+
@test sk == [1 3 5; 2 4 6]
58+
@test named_axiskeys(sk) == (a=[:x, :y], b=10:12)
59+
end
4960

5061
@test axiskeys(selectdim(M, :r, [true, false, true])) == (['a', 'c'], 2:5)
5162

test/_packages.jl

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -129,25 +129,25 @@ end
129129
end
130130
end
131131
@testset "stack" begin
132-
using LazyStack
132+
using LazyStack: stack as lstack
133133

134134
rin = [wrapdims(1:3, a='a':'c') for i=1:4]
135135

136-
@test axiskeys(stack(rin), :a) == 'a':'c'
137-
@test axiskeys(stack(:b, rin...), :a) == 'a':'c' # tuple
138-
@test axiskeys(stack(z for z in rin), :a) == 'a':'c' # generator
136+
@test axiskeys(lstack(rin), :a) == 'a':'c'
137+
@test axiskeys(lstack(:b, rin...), :a) == 'a':'c' # tuple
138+
@test axiskeys(lstack(z for z in rin), :a) == 'a':'c' # generator
139139

140140
rout = wrapdims([[1,2], [3,4]], b=10:11)
141-
@test axiskeys(stack(rout), :b) == 10:11
141+
@test axiskeys(lstack(rout), :b) == 10:11
142142

143143
rboth = wrapdims(rin, b=10:13)
144-
@test axiskeys(stack(rboth), :a) == 'a':'c'
145-
@test axiskeys(stack(rboth), :b) == 10:13
144+
@test axiskeys(lstack(rboth), :a) == 'a':'c'
145+
@test axiskeys(lstack(rboth), :b) == 10:13
146146

147147
nts = [(i=i, j="j", k=33) for i=1:3]
148-
@test axiskeys(stack(nts), 1) == [:i, :j, :k]
149-
@test axiskeys(stack(:z, nts...), 1) == [:i, :j, :k]
150-
@test axiskeys(stack(n for n in nts), 1) == [:i, :j, :k]
148+
@test axiskeys(lstack(nts), 1) == [:i, :j, :k]
149+
@test axiskeys(lstack(:z, nts...), 1) == [:i, :j, :k]
150+
@test axiskeys(lstack(n for n in nts), 1) == [:i, :j, :k]
151151

152152
end
153153
@testset "dates" begin

test/runtests.jl

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
using Test, AxisKeys, NamedDims
2-
using Statistics, OffsetArrays, Tables, UniqueVectors, LazyStack
2+
using Statistics, OffsetArrays, Tables, UniqueVectors
33
using ChainRulesCore: ProjectTo, NoTangent
44
using ChainRulesTestUtils: test_rrule
55
using FiniteDifferences

0 commit comments

Comments
 (0)