Index trait for resulting array's shape (#314)

JuliaArrays · Jun 13, 2022 · 1fb77f5 · 1fb77f5
1 parent 7c7daec
commit 1fb77f5
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 2 deletions.
diff --git a/lib/ArrayInterfaceCore/Project.toml b/lib/ArrayInterfaceCore/Project.toml
@@ -1,6 +1,6 @@
 name = "ArrayInterfaceCore"
 uuid = "30b0a656-2188-435a-8636-2ec0e6a096e2"
-version = "0.1.10"
+version = "0.1.11"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

diff --git a/lib/ArrayInterfaceCore/src/ArrayInterfaceCore.jl b/lib/ArrayInterfaceCore/src/ArrayInterfaceCore.jl
@@ -548,6 +548,20 @@ ndims_index(@nospecialize T::Type{<:Base.LogicalIndex}) = ndims(fieldtype(T, :ma
 ndims_index(T::Type) = 1
 ndims_index(@nospecialize(i)) = ndims_index(typeof(i))
 
+"""
+    ndims_shape(::Type{I}) -> Union{Int,Tuple{Vararg{Int}}}
+
+Returns the number of dimension that are represented in shape of the returned array when
+indexing with an instance of `I`.
+"""
+ndims_shape(T::DataType) = ndims_index(T)
+ndims_shape(::Type{Colon}) = 1
+ndims_shape(T::Type{<:Base.AbstractCartesianIndex{N}}) where {N} = ntuple(zero, Val{N}())
+ndims_shape(@nospecialize T::Type{<:CartesianIndices}) = ntuple(one, Val{ndims(T)}())
+ndims_shape(@nospecialize T::Type{<:Number}) = 0
+ndims_shape(@nospecialize T::Type{<:AbstractArray}) = ndims(T)
+ndims_shape(x) = ndims_shape(typeof(x))
+
 """
     instances_do_not_alias(::Type{T}) -> Bool
 

diff --git a/lib/ArrayInterfaceCore/test/runtests.jl b/lib/ArrayInterfaceCore/test/runtests.jl
@@ -268,6 +268,14 @@ end
     @test @inferred(ArrayInterfaceCore.ndims_index(1)) == 1
 end
 
+@testset "ndims_shape" begin
+    @test @inferred(ArrayInterfaceCore.ndims_shape(1)) === 0
+    @test @inferred(ArrayInterfaceCore.ndims_shape(:)) === 1
+    @test @inferred(ArrayInterfaceCore.ndims_shape(CartesianIndex(1, 2))) === (0, 0)
+    @test @inferred(ArrayInterfaceCore.ndims_shape(CartesianIndices((2,2)))) === (1, 1)
+    @test @inferred(ArrayInterfaceCore.ndims_shape([1 1])) === 2
+end
+
 @testset "indices_do_not_alias" begin
   @test ArrayInterfaceCore.instances_do_not_alias(Float64)
   @test !ArrayInterfaceCore.instances_do_not_alias(Matrix{Float64})
@@ -285,4 +293,3 @@ end
   @test !ArrayInterfaceCore.indices_do_not_alias(typeof(view(fill(rand(4,4),4,4)', 2:3, 1:2)))
   @test !ArrayInterfaceCore.indices_do_not_alias(typeof(view(rand(4,4)', StepRangeLen(1,0,5), 1:2)))
 end
-