Skip to content

Commit

Permalink
make axes(A) return KeyedUnitRanges
Browse files Browse the repository at this point in the history
  • Loading branch information
Michael Abbott authored and mcabbott committed Oct 6, 2021
1 parent f822f26 commit ff8824b
Show file tree
Hide file tree
Showing 3 changed files with 55 additions and 1 deletion.
2 changes: 2 additions & 0 deletions src/AxisKeys.jl
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,8 @@ module AxisKeys
include("struct.jl")
export KeyedArray, axiskeys

include("axes.jl")

include("lookup.jl")

include("names.jl")
Expand Down
52 changes: 52 additions & 0 deletions src/axes.jl
Original file line number Diff line number Diff line change
@@ -0,0 +1,52 @@
#=
This is an experiment, with returning something very similar to a KeyedArray from axes(A).
The reason to do so is that things like similar(C, axes(A,1), axes(B,2)) could propagate keys.
However right now axes(A,1) !isa Base.OneTo results in lots of OffsetArrays...
=#

struct KeyedUnitRange{T,AT,KT} <: AbstractUnitRange{T}
data::AT
keys::KT
function KeyedUnitRange(data::AT, keys::KT) where {AT<:AbstractUnitRange{T}, KT<:AbstractVector} where {T}
new{T,AT,KT}(data, keys)
end
end

Base.parent(A::KeyedUnitRange) = getfield(A, :data)
keyless(A::KeyedUnitRange) = parent(A)

for f in [:size, :first, :last, :IndexStyle]
@eval Base.$f(A::KeyedUnitRange) = $f(parent(A))
end

Base.getindex(A::KeyedUnitRange, inds::Integer) = getindex(parent(A), inds)

axiskeys(A::KeyedUnitRange) = tuple(getfield(A, :keys))
axiskeys(A::KeyedUnitRange, d::Integer) = d==1 ? getfield(A, :keys) : Base.OneTo(1)


# getkey(A::AbstractKeyedArray{<:Any,1}, key) = findfirst(isequal(key), axiskeys(A)[1])

function Base.axes(A::KeyedArray)
ntuple(ndims(A)) do d
KeyedUnitRange(axes(parent(A),d), axiskeys(A, d))
end
end

KeyedUnion{T} = Union{KeyedArray{T}, KeyedUnitRange{T}}

Base.summary(io::IO, x::KeyedUnion) = _summary(io, x)
Base.summary(io::IO, A::NamedDimsArray{L,T,N,<:KeyedUnion}) where {L,T,N} = _summary(io, A)
showtype(io::IO, ::KeyedUnitRange) = print(io, "KeyedUnitRange(...)")

function Base.show(io::IO, m::MIME"text/plain", x::KeyedUnitRange)
summary(io, x)
println(io, ":")
keyed_print_matrix(io, x)
end

function Base.show(io::IO, x::KeyedUnitRange)
print(io, "KeyedUnitRange(", keyless(x), ", ", axiskeys(x,1), ")")
end
2 changes: 1 addition & 1 deletion src/struct.jl
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@ KeyedArray(A::KeyedVector, k2::Tuple{AbstractVector}) =

Base.size(x::KeyedArray) = size(parent(x))

Base.axes(x::KeyedArray) = axes(parent(x))
# Base.axes(x::KeyedArray) = axes(parent(x))

Base.parent(x::KeyedArray) = getfield(x, :data)
keyless(x::KeyedArray) = parent(x)
Expand Down

0 comments on commit ff8824b

Please sign in to comment.