add tests, docstrings, fixes

Author: AntonOresten

src/Rewrap.jl

@@ -35,7 +35,7 @@ export Merge
 export Split, Split1
 export Resqueeze, Squeeze, Unsqueeze
 
-include("enhanced-base")
+include("enhanced-base/enhanced-base.jl")
 
 include("Permute.jl")
 export Permute

src/enhanced-base/dropdims.jl

@@ -1,5 +1,29 @@
+"""
+    dropdims(x; dims)
+
+Drop the specified dimensions from the array `x`.
+
+```jldoctest
+julia> x = [1 3 5; 2 4 6;;;]
+2×3×1 Array{Int64, 3}:
+[:, :, 1] =
+ 1  3  5
+ 2  4  6
+
+julia> y = view(x, :, 1:2, :)
+2×2×1 view(::Array{Int64, 3}, :, 1:2, :) with eltype Int64:
+[:, :, 1] =
+ 1  3
+ 2  4
+
+julia> Rewrap.dropdims(y; dims=3)
+2×2 view(::Matrix{Int64}, :, 1:2) with eltype Int64:
+ 1  3
+ 2  4
+```
+"""
 @constprop function dropdims(
-    x::AbstractArray{<:Any,N}; dims
+    x::AbstractArray{<:Any,N}; dims::Union{Int,Tuple{Vararg{Int}}}
 ) where N
     dims′ = dims isa Int ? (dims,) : dims
     ops = ntuple(i -> i in dims′ ? Squeeze() : Keep(), N)

src/enhanced-base/vec.jl

@@ -4,12 +4,23 @@
 Flatten the array `x` into a vector.
 
 ```jldoctest
-julia> x = rand(2, 3, 4);
+julia> x = [1 3 5; 2 4 6];
 
-julia> vec(x)
-24-element Vector{Float64}:
- 0.560475
- 0.188602
+julia> Rewrap.vec(view(x, :, 1:2))
+4-element view(::Vector{Int64}, 1:4) with eltype Int64:
+ 1
+ 2
+ 3
+ 4
+
+julia> Rewrap.vec(view(x, 1:2, :)) # not contiguous!
+6-element reshape(view(::Matrix{Int64}, 1:2, :), 6) with eltype Int64:
+ 1
+ 2
+ 3
+ 4
+ 5
+ 6
 ```
 """
 vec(x) = reshape(x, Merge(..))

test/runtests.jl

@@ -9,4 +9,5 @@ include("utils.jl")
     include("test_permute.jl")
     include("test_repeat.jl")
     include("test_reduce.jl")
+    include("test_enhanced_base.jl")
 end

test/test_enhanced_base.jl

@@ -0,0 +1,41 @@
+@testset "Enhanced Base" begin
+
+    @testset "dropdims" begin
+        A = reshape(collect(1:24), 4, 1, 3, 1, 2)
+
+        @test Rewrap.dropdims(A; dims=2) == reshape(A, 4, 3, 1, 2)
+        @test Rewrap.dropdims(A; dims=4) == reshape(A, 4, 1, 3, 2)
+        @test Rewrap.dropdims(A; dims=(2, 4)) == reshape(A, 4, 3, 2)
+        @test Rewrap.dropdims(A; dims=(4, 2)) == reshape(A, 4, 3, 2)
+
+        B = reshape(collect(1:6), 1, 2, 3, 1)
+        @test Rewrap.dropdims(B; dims=(1, 4)) == reshape(B, 2, 3)
+
+        x = view(reshape(collect(1:12), 4, 1, 3), :, 1:1, :)
+        y = Rewrap.dropdims(x; dims=2)
+        @test y == reshape(x, 4, 3)
+        @test _shares_storage(y, parent(x))
+    end
+
+    @testset "vec" begin
+        A = [1 3 5; 2 4 6]
+        @test Rewrap.vec(A) == [1, 2, 3, 4, 5, 6]
+
+        x = view(A, :, 1:2)
+        y = Rewrap.vec(x)
+        @test y == [1, 2, 3, 4]
+        @test y isa SubArray
+        @test _shares_storage(y, parent(x))
+
+        x2 = view(A, 1:2, :)
+        y2 = Rewrap.vec(x2)
+        @test y2 isa Base.ReshapedArray
+        @test y2 == [1, 2, 3, 4, 5, 6]
+
+        B = reshape(collect(1:24), 2, 3, 4)
+        @test Rewrap.vec(B) == collect(1:24)
+        @test Rewrap.vec(B) isa Array
+    end
+
+end
+