ENH: .squeeze has gained the axis parameter

closes #15339

Author: Kernc <[email protected]>

Closes #15335 from kernc/squeeze_axis_param and squashes the following commits:

44d3c54 [Kernc] fixup! ENH: .squeeze accepts axis parameter
cc018c9 [Kernc] ENH: .squeeze accepts axis parameter

Author: kernc

doc/source/whatsnew/v0.20.0.txt

@@ -149,6 +149,7 @@ Other enhancements
 - ``Series/DataFrame.asfreq()`` have gained a ``fill_value`` parameter, to fill missing values (:issue:`3715`).
 - ``Series/DataFrame.resample.asfreq`` have gained a ``fill_value`` parameter, to fill missing values during resampling (:issue:`3715`).
 - ``pandas.tools.hashing`` has gained a ``hash_tuples`` routine, and ``hash_pandas_object`` has gained the ability to hash a ``MultiIndex`` (:issue:`15224`)
+- ``Series/DataFrame.squeeze()`` have gained the ``axis`` parameter. (:issue:`15339`)
 
 .. _ISO 8601 duration: https://en.wikipedia.org/wiki/ISO_8601#Durations
 

pandas/compat/numpy/function.py

@@ -214,13 +214,6 @@ def validate_cum_func_with_skipna(skipna, args, kwargs, name):
 validate_stat_ddof_func = CompatValidator(STAT_DDOF_FUNC_DEFAULTS,
                                           method='kwargs')
 
-# Currently, numpy (v1.11) has backwards compatibility checks
-# in place so that this 'kwargs' parameter is technically
-# unnecessary, but in the long-run, this will be needed.
-SQUEEZE_DEFAULTS = dict(axis=None)
-validate_squeeze = CompatValidator(SQUEEZE_DEFAULTS, fname='squeeze',
-                                   method='kwargs')
-
 TAKE_DEFAULTS = OrderedDict()
 TAKE_DEFAULTS['out'] = None
 TAKE_DEFAULTS['mode'] = 'raise'

pandas/core/generic.py

@@ -532,13 +532,27 @@ def pop(self, item):
 
         return result
 
-    def squeeze(self, **kwargs):
-        """Squeeze length 1 dimensions."""
-        nv.validate_squeeze(tuple(), kwargs)
+    def squeeze(self, axis=None):
+        """
+        Squeeze length 1 dimensions.
 
+        Parameters
+        ----------
+        axis : None, integer or string axis name, optional
+            The axis to squeeze if 1-sized.
+
+            .. versionadded:: 0.20.0
+
+        Returns
+        -------
+        scalar if 1-sized, else original object
+        """
+        axis = (self._AXIS_NAMES if axis is None else
+                (self._get_axis_number(axis),))
         try:
-            return self.iloc[tuple([0 if len(a) == 1 else slice(None)
-                                    for a in self.axes])]
+            return self.iloc[
+                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

@@ -1770,17 +1770,27 @@ def test_squeeze(self):
         [tm.assert_series_equal(empty_series, higher_dim.squeeze())
          for higher_dim in [empty_series, empty_frame, empty_panel]]
 
+        # axis argument
+        df = tm.makeTimeDataFrame(nper=1).iloc[:, :1]
+        tm.assert_equal(df.shape, (1, 1))
+        tm.assert_series_equal(df.squeeze(axis=0), df.iloc[0])
+        tm.assert_series_equal(df.squeeze(axis='index'), df.iloc[0])
+        tm.assert_series_equal(df.squeeze(axis=1), df.iloc[:, 0])
+        tm.assert_series_equal(df.squeeze(axis='columns'), df.iloc[:, 0])
+        tm.assert_equal(df.squeeze(), df.iloc[0, 0])
+        tm.assertRaises(ValueError, df.squeeze, axis=2)
+        tm.assertRaises(ValueError, df.squeeze, axis='x')
+
+        df = tm.makeTimeDataFrame(3)
+        tm.assert_frame_equal(df.squeeze(axis=0), df)
+
     def test_numpy_squeeze(self):
         s = tm.makeFloatSeries()
         tm.assert_series_equal(np.squeeze(s), s)
 
         df = tm.makeTimeDataFrame().reindex(columns=['A'])
         tm.assert_series_equal(np.squeeze(df), df['A'])
 
-        msg = "the 'axis' parameter is not supported"
-        tm.assertRaisesRegexp(ValueError, msg,
-                              np.squeeze, s, axis=0)
-
     def test_transpose(self):
         msg = (r"transpose\(\) got multiple values for "
                r"keyword argument 'axes'")