diff --git a/src/Wrap/PyArray.jl b/src/Wrap/PyArray.jl index 64bb3176..95c639f1 100644 --- a/src/Wrap/PyArray.jl +++ b/src/Wrap/PyArray.jl @@ -639,7 +639,13 @@ function pyarray_offset(x::PyArray{T,N,F}, i::Int) where {T,N,F} end function pyarray_offset(x::PyArray{T,N,F}, i::Vararg{Int,N}) where {T,N,F} - sum((i .- 1) .* x.strides) + # Strides from the size make the unit first stride a compile-time constant, so loops vectorise. + # Also matters for N == 1, where this method is picked over `pyarray_offset(x, i::Int)`. + if (:contiguous in F) + sum((i .- 1) .* Utils.size_to_fstrides(sizeof(pyarray_get_R(T)), x.size)) + else + sum((i .- 1) .* x.strides) + end end function pyarray_load(::Type{T}, p::Ptr{R}) where {T,R} diff --git a/test/Wrap.jl b/test/Wrap.jl index 4720cdd4..38c7a40a 100644 --- a/test/Wrap.jl +++ b/test/Wrap.jl @@ -100,6 +100,18 @@ ) @test_throws Exception PyArray(nd; array = false, buffer = true) end + @testset "contiguous 2D" begin + tb = pyimport("_testbuffer") + nd = tb.ndarray( + pylist(1:6), + shape = pylist([2, 3]), + format = "i", + flags = tb.ND_FORTRAN, + ) + a = PyArray(nd; array = false) + @test a isa PyArray{Cint,2,(:linear,:contiguous)} + @test [a[i, j] for i = 1:2, j = 1:3] == reshape(1:6, 2, 3) + end end @testitem "PyDict" begin