allow mixed cartesian and regular indexing

rafaqz · Jun 1, 2022 · b5ff9b9 · b5ff9b9
1 parent 099d892
commit b5ff9b9
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 2 deletions.
diff --git a/src/array/indexing.jl b/src/array/indexing.jl
@@ -27,11 +27,18 @@ for f in (:getindex, :view, :dotview)
         @propagate_inbounds Base.$f(A::AbstractDimArray, args::Dimension...; kw...) =
             Base.$f(A, dims2indices(A, (args..., kwdims(values(kw))...))...)
         # Standard indices
-        @propagate_inbounds Base.$f(A::AbstractDimArray, i1::StandardIndices, i2::StandardIndices, I::StandardIndices...) =
-            rebuildsliced(Base.$f, A, Base.$f(parent(A), i1, i2, I...), (i1, i2, I...))
+        @propagate_inbounds function Base.$f(A::AbstractDimArray, i1::StandardIndices, i2::StandardIndices, I::StandardIndices...)
+            I = _unwrap_cartesian(i1, i2, I...)
+            x = Base.$f(parent(A), I...)
+            x isa AbstractArray ? rebuildsliced(Base.$f, A, x, I) : x
+        end
     end
 end
 
+@inline _unwrap_cartesian(i1::CartesianIndex, I...) = (Tuple(i1)..., _unwrap_cartesian(I...)...)
+@inline _unwrap_cartesian(i1, I...) = (i1, _unwrap_cartesian(I...)...)
+@inline _unwrap_cartesian() = ()
+
 #### setindex ####
 
 @propagate_inbounds Base.setindex!(A::AbstractDimArray, x) = setindex!(parent(A), x)

diff --git a/test/indexing.jl b/test/indexing.jl
@@ -65,6 +65,13 @@ end
         @test x = da[1, :][1:2] isa DimArray
     end
 
+    @testset "mixed CartesianIndex indexing works" begin
+
+        da3 = cat(da, 10da; dims=Z) 
+        @test da3[1, CartesianIndex(1, 2)] == 10
+        @test view(da3, 1:2, CartesianIndex(1, 2)) == [10, 30]
+    end
+
     @testset "getindex returns DimensionArray slices with the right dimensions" begin
         a = da[X(1:2), Y(1)]
         @test a == [1, 3]