Ideas for unifying array naming conventions in Julia.

named_array_traits.jl

import Base: tail, to_dim

struct Dim{Sym} end

const TimeDim = Dim{:time}()
const ColorDim = Dim{:color}()

struct HasDimNames{T} end
HasDimNames(x::T) where T = HasDimNames(T)
HasDimNames(::Type{T}) where T = HasDimNames{false}()

struct HasAxes{T} end
HasAxes(x::T) where T = HasAxes(T)
HasAxes(::Type{T}) where T = HasAxes{false}()
HasAxes(::Type{T}) where T<:AbstractArray = HasAxes{true}()

# - act like to `to_index` but for dimension names
# - adapted from `dim` in `NamedDims/named_core.jl`
"""
    to_dims(x, n)
"""
to_dims(x::Any, n::NTuple{N}) where {N} = map(to_dims, n)::NTuple{N,Int}

to_dims(x::T, n::Int) where {T} = _to_dims(HasAxes(T), n)
to_dims(x::T, n::Union{Symbol,Dim})::Int where {T} = _to_dims(HasDimNames(T), x, n)

function _to_dims(::HasDimNames{true}, x::Any, n::Symbol)
    dimnum = _to_dims_symbol(dimnames(x), n)
    if dimnum === 0
        throw(ArgumentError(
            "Specified name ($(repr(n))) does not match any dimension name ($dimnames(x))"
        ))
    end
    return dimnum
end

function _to_dims(::HasDimNames{true}, x::Any, n::Dim)
    dimnum = _to_dims_dim(dimnames(x), n)
    if dimnum === 0
        throw(ArgumentError(
            "Specified name ($(repr(n))) does not match any dimension name ($dimnames(x))"
        ))
    end
    return dimnum
end


_to_dims(::HasAxes{true}, x, n::Int) = _to_dims_int(axes(x), n)

Base.@pure function _to_dims_int(::Tuple{Vararg{Any,N}}, i::Int) where {N}
    if i < 1 || i > N
        throw(ArgumentError("Specified dimension, $i, not within $nd dimensions."))
    else
        return i
    end
end

Base.@pure function _to_dims_symbol(dimnames::NTuple{N,Symbol}, name::Symbol) where N
    for ii in 1:N
        getfield(dimnames, ii) === name && return ii
    end
    return 0
end

Base.@pure function _to_dims_dim(dimnames::NTuple{N,Symbol}, ::Dim{name}) where {N,name}
    for ii in 1:N
        getfield(dimnames, ii) === name && return ii
    end
    return 0
end

"""
    to_axes(x, n)
"""
to_axes(x::Any, n::NTuple{N}) = map(to_axes, n)

to_axes(x::T, i::Any) where {T} = _to_axes(HasAxes(T), x, i)
_to_axes(::HasAxes{true}, x, i::Int) = axes(x, i)
_to_axes(::HasAxes{true}, x, i::Symbol) = axes(x, to_dims(x, i))
_to_axes(::HasAxes{fasle}, x, i) = error("Type $(typeof(x)) does not have axes.")


"""
    namedaxes
"""
namedaxes(x::T) where T = _namedaxes(HasDimNames(T), HasAxes(T), x)
_namedaxes(::HasDimNames{true}, ::HasAxes{true}, x) = NamedTuple{dimnames(x)}(axes(x))
_namedaxes(::HasDimNames{Any}, ::HasAxes{Any}, x) = nothing

namedaxes(x::T, i) where {T} = _namedaxes(HasDimNames(T), HasAxes(T), x, i)
_namedaxes(::HasDimNames{true}, ::HasAxes{true}, x, i::Int) = NamedTuple{(dimnames(x, i),)}((axes(x, i),))
_namedaxes(::HasDimNames{true}, ::HasAxes{true}, x, i::Symbol) = NamedTuple{(i,)}((to_axes(x, i),))


# like `first` but also keeps the key for a NamedTuple
namedfirst(x::NamedTuple) = NamedTuple{(first(keys(x)),)}((first(x),))

"""
    filteraxes(f, x)
"""
filteraxes(f, x) = _catch_empty(_filteraxes(HasDimNames(x), f, x))
_filteraxes(::HasDimNames{true}, f, x) = __filteraxes(f, namedaxes(x))
_filteraxes(::HasDimNames{false}, f, x) = _findaxes(HasAxes(x), f, x)
_filteraxes(::HasAxes{true}, f, x) = __filteraxes(f, axes(x))
_filteraxes(::HasAxes{false}, f, x) = error("Type $(typeof(x)) does not have axes.")

function __filteraxes(f, t::Tuple)
    if f(first(t))
        return (first(t), __filteraxes(f, tail(t))...)
    else
        return __filteraxes(f, tail(t))
    end
end
function __filteraxes(f, t::NamedTuple)
    if f(namedfirst(t))
        return (namedfirst(t), __filteraxes(f, tail(t))...)
    else
        return __filteraxes(f, tail(t))
    end
end
_filteraxes(f, ::NamedTuple{(),Tuple{}}) = NamedTuple{(),Tuple{}}()
_filteraxes(f, ::Tuple{}) = ()

"""
    findaxes(f, x)
"""
findaxes(f, x) = _catch_empty(_findaxes(HasDimNames(x), f, x))
_findaxes(::HasDimNames{true}, f, x) = __findraxes(f, namedaxes(x), 1)
_findaxes(::HasDimNames{false}, f, x) = _findaxes(HasAxes(x), f, x)
_findaxes(::HasAxes{true}, f, x) = __findaxes(f, axes(x), 1)
_findaxes(::HasAxes{false}, f, x) = error("Type $(typeof(x)) does not have axes.")
function __findaxes(f, t::Tuple, cnt::Int)
    if f(first(t))
        return (cnt, __findaxes(f, tail(t), cnt+1)...)
    else
        return __findaxes(f, tail(t), cnt+1)
    end
end

function __findaxes(f, t::NamedTuple, cnt::Int)
    if f(namedfirst(t))
        return (cnt, __findaxes(f, tail(t), cnt+1)...)
    else
        return __findaxes(f, tail(t), cnt+1)
    end
end

__findaxes(f, ::Tuple{}) = ()
__findaxes(f, ::NamedTuple{(),Tuple{}}) = NamedTuple{(),Tuple{}}()

_catch_empty(x::Tuple) = x
_catch_empty(x::NamedTuple) = x
_catch_empty(::Tuple{}) = nothing
_catch_empty(::NamedTuple{(),Tuple{}}) = nothing