restrict: only return OffsetArray when input is an OffsetArray (#4)

Author: johnnychen94

src/ImageUtils.jl

@@ -4,7 +4,9 @@ export restrict
 
 using Base.Cartesian: @nloops
 using ImageCore
+using ImageCore.OffsetArrays
 
 include("restrict.jl")
+include("compat.jl")
 
 end

src/compat.jl

@@ -0,0 +1,5 @@
+if VERSION < v"1.2"
+    require_one_based_indexing(A...) = !Base.has_offset_axes(A...) || throw(ArgumentError("offset arrays are not supported but got an array with index other than 1"))
+else
+    const require_one_based_indexing = Base.require_one_based_indexing
+end

src/restrict.jl

@@ -76,9 +76,20 @@ function restrict(A::AbstractArray, region::Dims)
 end
 
 function restrict(A::AbstractArray{T,N}, dim::Integer) where {T,N}
+    require_one_based_indexing(A)
+
+    indsA = axes(A)
+    newinds = ntuple(i->i==dim ? restrict_indices(indsA[dim]) : indsA[i], Val(N))
+    out = Array{restrict_eltype(first(A)), N}(undef, last.(newinds))
+    restrict!(out, A, dim)
+    out
+end
+function restrict(A::OffsetArray{T,N}, dim::Integer) where {T,N}
     indsA = axes(A)
     newinds = map(UnitRange, ntuple(i->i==dim ? restrict_indices(indsA[dim]) : indsA[i], Val(N)))
-    out = similar(Array{restrict_eltype(first(A)), N}, newinds)
+    # This calls OffsetArrays implementation: a type piracy
+    # https://github.com/JuliaArrays/OffsetArrays.jl/issues/87
+    out = similar(A, restrict_eltype(first(A)), newinds)
     restrict!(out, A, dim)
     out
 end

test/restrict.jl

@@ -1,4 +1,23 @@
-@testset "restrict" begin   
+@testset "restrict" begin
+    @testset "interfaces" begin
+        A = rand(N0f8, 4, 5, 3)
+
+        Ar = @inferred restrict(A, 1)
+        @test typeof(Ar) <: Array
+        @test size(Ar) == (3, 5, 3)
+
+        Ar = @inferred restrict(A, (1, ))
+        @test typeof(Ar) <: Array
+        @test size(Ar) == (3, 5, 3)
+
+        Ar = @inferred restrict(A, (1, 2, 3))
+        @test typeof(Ar) <: Array
+        @test size(Ar) == (3, 3, 2)
+        @test Ar == restrict(A)
+
+        @test_throws MethodError restrict(A, 1, 2, 3)
+    end
+
     @testset "numerical test" begin
         A = reshape([UInt16(i) for i = 1:60], 4, 5, 3)
         B = restrict(A, (1,2))
@@ -36,9 +55,17 @@
     @testset "OffsetArray" begin
         A = rand(5, 4, 3)
         Ao = OffsetArray(A, (-2,1,0))
-        @test parent(@inferred(restrict(Ao, 1))) == restrict(A, 1)
-        @test parent(@inferred(restrict(Ao, 2))) == restrict(A, 2)
-        @test parent(@inferred(restrict(Ao, (1,2)))) == restrict(A, (1,2))
+
+        for (dims, offsets) in [
+                (1,      (-1, 1, 0)),
+                (2,      (-2, 0, 0)),
+                ((1, 2), (-1, 0, 0))
+            ]
+            Ar = @inferred restrict(Ao, dims)
+            @test typeof(Ar) <: OffsetArray
+            @test Ar.offsets == offsets
+            @test parent(Ar) == restrict(A, dims)
+        end
     end
 
     @testset "FixedPoint overflow" begin