diff --git a/src/AxisArrays.jl b/src/AxisArrays.jl index 2d023cb..04955a4 100644 --- a/src/AxisArrays.jl +++ b/src/AxisArrays.jl @@ -7,7 +7,7 @@ import Base.Iterators: repeated using RangeArrays, IntervalSets using IterTools -export AxisArray, Axis, axisnames, axisvalues, axisdim, axes, atindex, atvalue, collapse +export AxisArray, Axis, axisnames, axisvalues, axisdim, axes, atindex, atvalue, collapse, @v # From IntervalSets: export ClosedInterval, .. diff --git a/src/indexing.jl b/src/indexing.jl index d8e9e7d..72f2c38 100644 --- a/src/indexing.jl +++ b/src/indexing.jl @@ -362,3 +362,55 @@ end # arrays within Base's to_index/to_indices methods, but that requires a bigger # refactor to merge our to_index method with Base's. @inline Base.checkindex(::Type{Bool}, inds::AbstractUnitRange, A::AxisArray) = Base.checkindex(Bool, inds, A.data) + +""" + @v expr +For a `getindex` or `setindex!` operation writen in bracket notation, wrap all indices +except for `:` such that `i → atvalue(i)`. You can avoid wrapping in `atvalue`, or manually +wrap so that you can specify `atol` and `rtol`, by splicing in an index as shown in the +examples. This behavior is analogous to that of the broadcasting `@.` macro. + +# Examples + +```jldoctest +julia> A = AxisArray(reshape(1:6, 2, 3), Axis{:a}([-1.1, 0]), Axis{:b}([:cat, :lemur, :olinguito])); +2-dimensional AxisArray{Int64,2,...} with axes: + :a, -1:0 + :b, Symbol[:cat, :lemur, :olinguito] +And data, a 2×3 Base.ReshapedArray{Int64,2,UnitRange{Int64},Tuple{}}: + 1 3 5 + 2 4 6 + +julia> @v A[-1.1, :cat] +1 + +julia> @v A[-1.1,\$2] +3 + +julia> @v A[\$(atvalue(-1.1, atol=0.2)), :olinguito] +5 +``` +""" +macro v(x::Expr) + esc(v(x)) +end + +v(x) = x +function v(x::Expr) + if x.head == :ref + for i in 2:length(x.args) + if x.args[i] isa Expr + if x.args[i].head == :$ + x.args[i] = x.args[i].args[1] + end + elseif x.args[i] != :(:) + x.args[i] = :(atvalue($(x.args[i]))) + end + end + else + for i in 1:length(x.args) + v(x.args[i]) + end + end + return x +end diff --git a/test/indexing.jl b/test/indexing.jl index 19ae91c..b102bee 100644 --- a/test/indexing.jl +++ b/test/indexing.jl @@ -283,3 +283,19 @@ A = AxisArray(1:365, Date(2017,1,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] + +# `@v` macro for indexing by value +## getindex behavior +A = AxisArray(collect(reshape(1:6, 2, 3)), + Axis{:a}([-1.1, 0]), Axis{:b}([:cat, :lemur, :olinguito])); +@test (@v A[-1.1, :cat]) == A[1,1] +@test (@v A[-1.1, $2]) == A[1,2] +@test (@v A[$(atvalue(-1.1, atol=0.2)), :olinguito]) == A[1,3] +@test (@v A[0, :]) == A[2,:] +@test (@v A[-0.5..0, :cat]) == AxisArray([2], Axis{:a}([0.0])) + +## setindex! behavior +@v A[-1.1, :cat] = 20 +@test @v A[-1.1, :cat] == 20 +@v A[:, :lemur] .= 12 +@test @v A[:, :lemur] == [12,12]