From 1fb77f5ee8f71a84202c9eb856856b96ef9ceb2f Mon Sep 17 00:00:00 2001
From: Zachary P Christensen <zchristensen7@gmail.com>
Date: Mon, 13 Jun 2022 13:51:23 -0400
Subject: [PATCH] Index trait for resulting array's shape (#314)

---
 lib/ArrayInterfaceCore/Project.toml              |  2 +-
 lib/ArrayInterfaceCore/src/ArrayInterfaceCore.jl | 14 ++++++++++++++
 lib/ArrayInterfaceCore/test/runtests.jl          |  9 ++++++++-
 3 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/lib/ArrayInterfaceCore/Project.toml b/lib/ArrayInterfaceCore/Project.toml
index c6f51d264..ec54f48e1 100644
--- 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
index 6d7d054fe..023194858 100644
--- 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
index 291883cf0..1e79b0341 100644
--- 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
-