Merge pull request #93 from Tokazama/master

Easier use of ArrayInterface.size/axes/to_dims

Author: chriselrod

Project.toml

@@ -1,6 +1,6 @@
 name = "ArrayInterface"
 uuid = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
-version = "2.14.2"
+version = "2.14.3"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

src/dimensions.jl

@@ -71,13 +71,14 @@ end
 
 This returns the dimension(s) of `x` corresponding to `d`.
 """
-to_dims(x, d::Integer) = Int(d)
-to_dims(x, d::Colon) = d   # `:` is the default for most methods that take `dims`
-@inline to_dims(x, d::Tuple) = map(i -> to_dims(x, i), d)
-@inline function to_dims(x, d::Symbol)::Int
-    i = _sym_to_dim(dimnames(x), d)
+to_dims(x, d) = to_dims(dimnames(x), d)
+to_dims(x::Tuple{Vararg{Symbol}}, d::Integer) = Int(d)
+to_dims(x::Tuple{Vararg{Symbol}}, d::Colon) = d   # `:` is the default for most methods that take `dims`
+@inline to_dims(x::Tuple{Vararg{Symbol}}, d::Tuple) = map(i -> to_dims(x, i), d)
+@inline function to_dims(x::Tuple{Vararg{Symbol}}, d::Symbol)::Int
+    i = _sym_to_dim(x, d)
     if i === 0
-        throw(ArgumentError("Specified name ($(repr(d))) does not match any dimension name ($(dimnames(x)))"))
+        throw(ArgumentError("Specified name ($(repr(d))) does not match any dimension name ($(x))"))
     end
     return i
 end
@@ -152,6 +153,7 @@ julia> ArrayInterface.size(A)
 ```
 """
 @inline size(A) = Base.size(A)
+@inline size(A, d::Integer) = size(A)[Int(d)]
 @inline size(A, d) = Base.size(A, to_dims(A, d))
 @inline size(x::LinearAlgebra.Adjoint{T,V}) where {T, V <: AbstractVector{T}} = (One(), static_length(x))
 @inline size(x::LinearAlgebra.Transpose{T,V}) where {T, V <: AbstractVector{T}} = (One(), static_length(x))
@@ -161,7 +163,8 @@ julia> ArrayInterface.size(A)
 
 Return a valid range that maps to each index along dimension `d` of `A`.
 """
-@inline axes(A, d) = Base.axes(A, to_dims(A, d))
+@inline axes(A, d) = axes(A, to_dims(A, d))
+@inline axes(A, d::Integer) = axes(A)[Int(d)]
 
 """
     axes(A)

test/dimensions.jl

@@ -1,6 +1,21 @@
 
 @testset "dimensions" begin
 
+@testset "to_dims" begin
+    @testset "small case" begin
+        @test ArrayInterface.to_dims((:x, :y), :x) == 1
+        @test ArrayInterface.to_dims((:x, :y), :y) == 2
+        @test_throws ArgumentError ArrayInterface.to_dims((:x, :y), :z)  # not found
+    end
+
+    @testset "large case" begin
+        @test ArrayInterface.to_dims((:x, :y, :a, :b, :c, :d), :x) == 1
+        @test ArrayInterface.to_dims((:x, :y, :a, :b, :c, :d), :a) == 3
+        @test ArrayInterface.to_dims((:x, :y, :a, :b, :c, :d), :d) == 6
+        @test_throws ArgumentError ArrayInterface.to_dims((:x, :y, :a, :b, :c, :d), :z) # not found
+    end
+end
+
 struct NamedDimsWrapper{L,T,N,P<:AbstractArray{T,N}} <: AbstractArray{T,N}
     parent::P
     NamedDimsWrapper{L}(p) where {L} = new{L,eltype(p),ndims(p),typeof(p)}(p)
@@ -10,8 +25,11 @@ ArrayInterface.has_dimnames(::Type{T}) where {T<:NamedDimsWrapper} = true
 ArrayInterface.dimnames(::Type{T}) where {L,T<:NamedDimsWrapper{L}} = L
 Base.parent(x::NamedDimsWrapper) = x.parent
 Base.size(x::NamedDimsWrapper) = size(parent(x))
+Base.size(x::NamedDimsWrapper, d) = ArrayInterface.size(x, d)
 Base.axes(x::NamedDimsWrapper) = axes(parent(x))
+Base.axes(x::NamedDimsWrapper, d) = ArrayInterface.axes(x, d)
 Base.strides(x::NamedDimsWrapper) = Base.strides(parent(x))
+Base.strides(x::NamedDimsWrapper, d) =  ArrayInterface.strides(x, d)
 
 Base.getindex(x::NamedDimsWrapper; kwargs...) = ArrayInterface.getindex(x; kwargs...)
 Base.getindex(x::NamedDimsWrapper, args...) = ArrayInterface.getindex(x, args...)
@@ -46,9 +64,9 @@ dnums = ntuple(+, length(d))
 @test @inferred(ArrayInterface.to_dims(x, reverse(d))) === reverse(dnums)
 @test_throws ArgumentError ArrayInterface.to_dims(x, :z)
 
-@test @inferred(ArrayInterface.size(x, :x)) == size(parent(x), 1)
-@test @inferred(ArrayInterface.axes(x, :x)) == axes(parent(x), 1)
-@test ArrayInterface.strides(x, :x) == ArrayInterface.strides(parent(x))[1]
+@test @inferred(size(x, :x)) == size(parent(x), 1)
+@test @inferred(axes(x, :x)) == axes(parent(x), 1)
+@test strides(x, :x) == ArrayInterface.strides(parent(x))[1]
 
 x[x = 1] = [2, 3]
 @test @inferred(getindex(x, x = 1)) == [2, 3]