Try converting indexes once and then defer to the parent

AxisArrays has to handle many different index types, but it can't
handle types defined in other packages. This provides for a fallback
mode in which AxisArrays tries to convert "fancy" index types into
"regular" ones, but if this fails then it assumes this is something
the parent knows how to handle.

Author: timholy

src/indexing.jl

@@ -114,13 +114,22 @@ end
 @propagate_inbounds Base.setindex!(A::AxisArray, v, idx::AbstractArray{Bool}) = (A.data[idx] = v)
 
 ### Fancier indexing capabilities provided only by AxisArrays ###
-@propagate_inbounds Base.getindex(A::AxisArray, idxs...) = A[to_index(A,idxs...)...]
-@propagate_inbounds Base.setindex!(A::AxisArray, v, idxs...) = (A[to_index(A,idxs...)...] = v)
+# To avoid StackOverflowErrors on indexes that we don't know how to convert, we
+# give AxisArrays once chance to convert into known format and then defer to the parent
+@propagate_inbounds Base.getindex(A::AxisArray, idxs...) = getindex_converted(A, to_index(A,idxs...)...)
+@propagate_inbounds Base.setindex!(A::AxisArray, v, idxs...) = setindex!_converted(A, v, to_index(A,idxs...)...)
 # Deal with lots of ambiguities here
 @propagate_inbounds Base.view(A::AxisArray, idxs::ViewIndex...) = view(A, to_index(A,idxs...)...)
 @propagate_inbounds Base.view(A::AxisArray, idxs::Union{ViewIndex,AbstractCartesianIndex}...) = view(A, to_index(A,Base.IteratorsMD.flatten(idxs)...)...)
 @propagate_inbounds Base.view(A::AxisArray, idxs...) = view(A, to_index(A,idxs...)...)
 
+@propagate_inbounds getindex_converted(A, idxs::Idx...) = A[idxs...]
+@propagate_inbounds setindex!_converted(A, v, idxs::Idx...) = (A[idxs...] = v)
+
+@propagate_inbounds getindex_converted(A, idxs...) = A.data[idxs...]
+@propagate_inbounds setindex!_converted(A, v, idxs...) = (A.data[idxs...] = v)
+
+
 # First is indexing by named axis. We simply sort the axes and re-dispatch.
 # When indexing by named axis the shapes of omitted dimensions are preserved
 # TODO: should we handle multidimensional Axis indexes? It could be interpreted
@@ -179,7 +188,7 @@ axisindexes(t, ax, idx) = error("cannot index $(typeof(ax)) with $(typeof(idx));
 # Maybe extend error message to all <: Numbers if Base allows it?
 axisindexes(::Type{Dimensional}, ax::AbstractVector, idx::Real) =
     throw(ArgumentError("invalid index: $idx. Use `atvalue` when indexing by value."))
-function axisindexes(::Type{Dimensional}, ax::AbstractVector, idx)
+function axisindexes(::Type{Dimensional}, ax::AbstractVector{T}, idx::T) where T
     idxs = searchsorted(ax, ClosedInterval(idx,idx))
     length(idxs) > 1 && error("more than one datapoint lies on axis value $idx; use an interval to return all values")
     if length(idxs) == 1
@@ -218,6 +227,8 @@ function axisindexes(::Type{Dimensional}, ax::AbstractVector, idx::ExactValue)
         throw(BoundsError(ax, idx))
     end
 end
+# For index types that AxisArrays doesn't know about
+axisindexes(::Type{Dimensional}, ax::AbstractVector, idx) = idx
 
 # Dimensional axes may be indexed by intervals to select a range
 axisindexes(::Type{Dimensional}, ax::AbstractVector, idx::ClosedInterval) = searchsorted(ax, idx)

test/indexing.jl

@@ -285,3 +285,20 @@ A = AxisArray(1:365, Date(2017,1,1):Day(1):Date(2017,12,31))
 @test A[(-Day(13)..Day(14)) + Date(2017,2,14)] == collect(31 .+ (1:28))
 @test A[(-Day(14)..Day(14)) + DateTime(2017,2,14,12)] == collect(31 .+ (1:28))
 @test A[(Day(0)..Day(6)) + (Date(2017,1,1):Month(1):Date(2017,4,12))] == [1:7 32:38 60:66 91:97]
+
+# Test using index types that AxisArrays doesn't understand
+# This example is inspired from Interpolations.jl and would implement linear interpolation
+struct WeightedIndex{T}
+    idx::Int
+    weights::Tuple{T,T}
+end
+Base.to_indices(A, I::Tuple{Vararg{Union{Int,WeightedIndex}}}) = I
+@inline function Base._getindex(::IndexStyle, A::AbstractVector, I::WeightedIndex)
+    A[I.idx]*I.weights[1] + A[I.idx+1]*I.weights[2]
+end
+idx = WeightedIndex(2, (0.2, 0.8))
+
+a = [2, 3, 7]
+@test a[idx] ≈ 6.2
+aa = AxisArray(a, :x)
+@test aa[idx] ≈ 6.2