From cc018c9c7f3f965e9356dee10a8eebefe6f244e2 Mon Sep 17 00:00:00 2001 From: Kernc Date: Tue, 7 Feb 2017 20:30:08 +0100 Subject: [PATCH 1/2] ENH: .squeeze accepts axis parameter --- doc/source/whatsnew/v0.20.0.txt | 2 ++ pandas/compat/numpy/function.py | 7 ------- pandas/core/generic.py | 35 ++++++++++++++++++++++++++++----- pandas/tests/test_generic.py | 16 +++++++++++---- 4 files changed, 44 insertions(+), 16 deletions(-) diff --git a/doc/source/whatsnew/v0.20.0.txt b/doc/source/whatsnew/v0.20.0.txt index 3f6c06e20b546..c1e8d35bbf1de 100644 --- a/doc/source/whatsnew/v0.20.0.txt +++ b/doc/source/whatsnew/v0.20.0.txt @@ -150,6 +150,8 @@ Other enhancements - ``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 support for ``axis`` parameter. (:issue:15339``) + .. _ISO 8601 duration: https://en.wikipedia.org/wiki/ISO_8601#Durations .. _whatsnew_0200.api_breaking: diff --git a/pandas/compat/numpy/function.py b/pandas/compat/numpy/function.py index 72e89586d0280..eb9e9ecc359b2 100644 --- a/pandas/compat/numpy/function.py +++ b/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' diff --git a/pandas/core/generic.py b/pandas/core/generic.py index bb2664a5b8d28..e736028edbf7f 100644 --- a/pandas/core/generic.py +++ b/pandas/core/generic.py @@ -532,13 +532,38 @@ 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 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. + + .. 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') try: - return self.iloc[tuple([0 if len(a) == 1 else slice(None) - for a in self.axes])] + return self.iloc[ + tuple([0 if len(a) == 1 and i in axis else slice(None) + for i, a in enumerate(self.axes)])] except: return self diff --git a/pandas/tests/test_generic.py b/pandas/tests/test_generic.py index 916d7ae0b0ec4..437408b6d6a06 100644 --- a/pandas/tests/test_generic.py +++ b/pandas/tests/test_generic.py @@ -1770,6 +1770,18 @@ 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=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') + + df = tm.makeTimeDataFrame(3) + tm.assertRaises(ValueError, df.squeeze, axis=0) + def test_numpy_squeeze(self): s = tm.makeFloatSeries() tm.assert_series_equal(np.squeeze(s), s) @@ -1777,10 +1789,6 @@ def test_numpy_squeeze(self): 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'") From 44d3c54edfc5927c584aab63b5e059dc8e594cb4 Mon Sep 17 00:00:00 2001 From: Kernc Date: Wed, 8 Feb 2017 02:11:24 +0100 Subject: [PATCH 2/2] fixup! ENH: .squeeze accepts axis parameter --- pandas/core/generic.py | 26 +++++--------------------- pandas/tests/test_generic.py | 4 ++-- 2 files changed, 7 insertions(+), 23 deletions(-) diff --git a/pandas/core/generic.py b/pandas/core/generic.py index e736028edbf7f..c9ebf14f165bc 100644 --- a/pandas/core/generic.py +++ b/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 diff --git a/pandas/tests/test_generic.py b/pandas/tests/test_generic.py index 437408b6d6a06..78fbdfa3d6798 100644 --- a/pandas/tests/test_generic.py +++ b/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()