diff --git a/src/array_partition.jl b/src/array_partition.jl index 2828b159..a1220451 100644 --- a/src/array_partition.jl +++ b/src/array_partition.jl @@ -317,12 +317,14 @@ Base.@propagate_inbounds function Base.getindex(A::ArrayPartition, i::Int) end """ - getindex(A::ArrayPartition, i::Colon, j...) + getindex(A::ArrayPartition, i::Colon, j, js...) -Returns the entry at index `j...` of every partition of `A`. +Returns the entry at index `j, js...` of every partition of `A`. +Requires at least one index after the colon so that `A[:]` keeps the +flattening `getindex(::ArrayPartition, ::Colon)` method. """ -Base.@propagate_inbounds function Base.getindex(A::ArrayPartition, i::Colon, j...) - return getindex.(A.x, (j...,)) +Base.@propagate_inbounds function Base.getindex(A::ArrayPartition, i::Colon, j, js...) + return getindex.(A.x, (j, js...)) end """ diff --git a/test/Core/partitions_test.jl b/test/Core/partitions_test.jl index 44bf10c9..93a3e7b8 100644 --- a/test/Core/partitions_test.jl +++ b/test/Core/partitions_test.jl @@ -90,6 +90,30 @@ x = AP[[1, 2], [3.0, 4.0]] y = ArrayPartition(AP[[1], [2.0]], AP[[3], [4.0]]) @test x[:, 1] == (1, 3.0) +# Issue #661: A[:] must flatten, not hit the (colon, j...) broadcast path +@testset "ArrayPartition A[:] flattening (issue #661)" begin + v = ArrayPartition([1.0, 2.0], [3.0]) + @test v[:] == [1.0, 2.0, 3.0] + @test v[:] == vec(v) + @test v[:] == collect(v) + @test v[1:3] == [1.0, 2.0, 3.0] + + m = ArrayPartition([1.0 2.0], [3.0 4.0 5.0]) + @test m[:] == [1.0, 2.0, 3.0, 4.0, 5.0] + @test m[:] == vec(m) + @test m[:] == collect(m) + @test m[1:4] == [1.0, 2.0, 3.0, 4.0] + + nested = ArrayPartition(ArrayPartition([1.0], [2.0]), ArrayPartition([3.0, 4.0])) + @test nested[:] == [1.0, 2.0, 3.0, 4.0] + @test nested[:] == vec(nested) + @test nested[:] == collect(nested) + + # (colon, j...) from #239 must still work + @test x[:, 1] == (1, 3.0) + @test x[:, 2] == (2, 4.0) +end + # similar partitions @inferred similar(x) @test similar(x, (4,)) isa ArrayPartition{Float64}