fixup! ENH: .squeeze accepts axis parameter

Author: kernc

pandas/core/generic.py

@@ -537,32 +537,16 @@ def squeeze(self, axis=None):
 
         Parameters
         ----------
-        axis : None or int or tuple of ints, optional
-            Selects a subset of the single-dimensional entries in the
-            shape. If an axis is selected with shape entry greater than
-            one, an error is raised.
+        axis : None or int, optional
+            The axis to squeeze if 1-sized.
 
             .. versionadded:: 0.20.0
         """
-        if axis is None:
-            axis = tuple(range(len(self.axes)))
-        else:
-            if not is_list_like(axis):
-                axis = (axis,)
-            if not all(is_integer(ax) for ax in axis):
-                raise TypeError('an integer is required for the axis')
-            n_axes = len(self.axes)
-            for ax in axis:
-                if ax < -n_axes or ax >= n_axes:
-                    raise ValueError("'axis' entry {0} is out of bounds "
-                                     "[-{1}, {1})".format(ax, n_axes))
-            if any(len(self.axes[ax]) != 1 for ax in axis):
-                raise ValueError('cannot select an axis to squeeze out which '
-                                 'has size not equal to one')
-
+        axis = (self._AXIS_NAMES if axis is None else
+                (self._get_axis_number(axis),))
         try:
             return self.iloc[
-                tuple([0 if len(a) == 1 and i in axis else slice(None)
+                tuple([0 if i in axis and len(a) == 1 else slice(None)
                        for i, a in enumerate(self.axes)])]
         except:
             return self

pandas/tests/test_generic.py

@@ -1777,10 +1777,10 @@ def test_squeeze(self):
         tm.assert_series_equal(df.squeeze(axis=1), df.iloc[:, 0])
         tm.assert_equal(df.squeeze(), df.iloc[0, 0])
         tm.assertRaises(ValueError, df.squeeze, axis=2)
-        tm.assertRaises(TypeError, df.squeeze, axis='x')
+        tm.assertRaises(ValueError, df.squeeze, axis='x')
 
         df = tm.makeTimeDataFrame(3)
-        tm.assertRaises(ValueError, df.squeeze, axis=0)
+        tm.assert_frame_equal(df.squeeze(axis=0), df)
 
     def test_numpy_squeeze(self):
         s = tm.makeFloatSeries()