Support squeeze

Author: timholy

src/AxisArrays.jl

@@ -1,5 +1,6 @@
 module AxisArrays
 
+using Base: tail
 using RangeArrays, Iterators, IntervalSets, Compat
 using Compat.view
 

src/core.jl

@@ -307,6 +307,23 @@ function permutation(to::Symbols, from::Symbols)
     ind
 end
 
+function Base.squeeze(A::AxisArray, dims::Dims)
+    keepdims = setdiff(1:ndims(A), dims)
+    AxisArray(squeeze(A.data, dims), axes(A)[keepdims])
+end
+# This version is type-stable
+function Base.squeeze{Ax<:Axis}(A::AxisArray, ::Type{Ax})
+    dim = axisdim(A, Ax)
+    AxisArray(squeeze(A.data, dim), dropax(Ax, axes(A)...))
+end
+
+@inline dropax(ax, ax1, axs...) = (ax1, dropax(ax, axs...)...)
+@inline dropax{name}(ax::Axis{name}, ax1::Axis{name}, axs...) = dropax(ax, axs...)
+@inline dropax{name}(ax::Type{Axis{name}}, ax1::Axis{name}, axs...) = dropax(ax, axs...)
+@inline dropax{name,T}(ax::Type{Axis{name,T}}, ax1::Axis{name}, axs...) = dropax(ax, axs...)
+dropax(ax) = ()
+
+
 # A simple display method to include axis information. It might be nice to
 # eventually display the axis labels alongside the data array, but that is
 # much more difficult.

test/core.jl

@@ -96,6 +96,14 @@ A = AxisArray(reshape(1:16, 2,2,2,2), .5:.5:1)
 @test axisnames(A) == (:row,:col,:page,:dim_4)
 VERSION >= v"0.5.0-dev" && @inferred(axisnames(A))
 @test axisvalues(A) == (.5:.5:1, 1:2, 1:2, 1:2)
+A = AxisArray([0]', :x, :y)
+@test axisnames(squeeze(A, 1)) == (:y,)
+@test axisnames(squeeze(A, 2)) == (:x,)
+@test axisnames(squeeze(A, (1,2))) == axisnames(squeeze(A, (2,1))) == ()
+@test axisnames(@inferred(squeeze(A, Axis{:x}))) == (:y,)
+@test axisnames(@inferred(squeeze(A, Axis{:x,UnitRange{Int}}))) == (:y,)
+@test axisnames(@inferred(squeeze(A, Axis{:y}))) == (:x,)
+@test axisnames(@inferred(squeeze(squeeze(A, Axis{:x}), Axis{:y}))) == ()
 
 # Test axisdim
 @test_throws ArgumentError AxisArray(reshape(1:24, 2,3,4),