BUG: Work around evil matrix.__array_prepare__

2 files changed, 15 insertions, 4 deletions

diff --git a/numpy/lib/shape_base.py b/numpy/lib/shape_base.py
index 7e70904b1..120153ad5 100644
--- a/numpy/lib/shape_base.py
+++ b/numpy/lib/shape_base.py
@@ -9,6 +9,7 @@ from numpy.core.numeric import (
 from numpy.core.fromnumeric import product, reshape, transpose
 from numpy.core import vstack, atleast_3d
 from numpy.lib.index_tricks import ndindex
+from numpy.matrixlib.defmatrix import matrix  # this raises all the right alarm bells
 
 
 __all__ = [
@@ -116,7 +117,11 @@ def apply_along_axis(func1d, axis, arr, *args, **kwargs):
     # insert as many axes as necessary to create the output
     outshape = arr.shape[:axis] + res.shape + arr.shape[axis+1:]
     outarr = zeros(outshape, res.dtype)
-    outarr = res.__array_prepare__(outarr)
+    # matrices call reshape in __array_prepare__, and this is apparently ok!
+    if not isinstance(res, matrix):
+        outarr = res.__array_prepare__(outarr)
+    if outshape != outarr.shape:
+        raise ValueError('__array_prepare__ should not change the shape of the resultant array')
 
     # outarr, with inserted dimensions at the end
     # this means that outarr_view[i] = func1d(inarr_view[i])
diff --git a/numpy/lib/tests/test_shape_base.py b/numpy/lib/tests/test_shape_base.py
index 7bf2b4a81..111f302aa 100644
--- a/numpy/lib/tests/test_shape_base.py
+++ b/numpy/lib/tests/test_shape_base.py
@@ -28,14 +28,20 @@ class TestApplyAlongAxis(TestCase):
                            [[27, 30, 33], [36, 39, 42], [45, 48, 51]])
 
     def test_preserve_subclass(self):
+        # this test is particularly malicious because matrix
+        # refuses to become 1d
         def double(row):
             return row * 2
         m = np.matrix([[0, 1], [2, 3]])
+        expected = np.matrix([[0, 2], [4, 6]])
+
         result = apply_along_axis(double, 0, m)
         assert_(isinstance(result, np.matrix))
-        assert_array_equal(
-            result, np.matrix([[0, 2], [4, 6]])
-        )
+        assert_array_equal(result, expected)
+
+        result = apply_along_axis(double, 1, m)
+        assert_(isinstance(result, np.matrix))
+        assert_array_equal(result, expected)
 
     def test_subclass(self):
         class MinimalSubclass(np.ndarray):