Merge pull request #164 from JuliaArrays/teh/unknown_index_types

Try converting indexes once and then defer to the parent

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