diff --git a/doc/source/whatsnew/v0.21.0.txt b/doc/source/whatsnew/v0.21.0.txt index 32dbeb32154e6..21abdccd2996c 100644 --- a/doc/source/whatsnew/v0.21.0.txt +++ b/doc/source/whatsnew/v0.21.0.txt @@ -91,6 +91,24 @@ This does not raise any obvious exceptions, but also does not create a new colum Setting a list-like data structure into a new attribute now raise a ``UserWarning`` about the potential for unexpected behavior. See :ref:`Attribute Access `. +``drop`` now also accepts index/columns keywords +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +The :meth:`~DataFrame.drop` method has gained ``index``/``columns`` keywords as an +alternative to specify the ``axis`` and to make it similar in usage to ``reindex`` +(:issue:`12392`). + +For example: + +.. ipython:: python + + df = pd.DataFrame(np.arange(8).reshape(2,4), + columns=['A', 'B', 'C', 'D']) + df + df.drop(['B', 'C'], axis=1) + # the following is now equivalent + df.drop(columns=['B', 'C']) + .. _whatsnew_0210.enhancements.categorical_dtype: ``CategoricalDtype`` for specifying categoricals diff --git a/pandas/core/generic.py b/pandas/core/generic.py index 241204ef555f6..3d55e07df6eac 100644 --- a/pandas/core/generic.py +++ b/pandas/core/generic.py @@ -2333,14 +2333,23 @@ def reindex_like(self, other, method=None, copy=True, limit=None, return self.reindex(**d) - def drop(self, labels, axis=0, level=None, inplace=False, errors='raise'): + def drop(self, labels=None, axis=0, index=None, columns=None, level=None, + inplace=False, errors='raise'): """ Return new object with labels in requested axis removed. Parameters ---------- labels : single label or list-like + Index or column labels to drop. axis : int or axis name + Whether to drop labels from the index (0 / 'index') or + columns (1 / 'columns'). + index, columns : single label or list-like + Alternative to specifying `axis` (``labels, axis=1`` is + equivalent to ``columns=labels``). + + .. versionadded:: 0.21.0 level : int or level name, default None For MultiIndex inplace : bool, default False @@ -2354,36 +2363,80 @@ def drop(self, labels, axis=0, level=None, inplace=False, errors='raise'): Examples -------- - >>> df = pd.DataFrame([[1, 2, 3, 4], - ... [5, 6, 7, 8], - ... [9, 1, 2, 3], - ... [4, 5, 6, 7] - ... ], - ... columns=list('ABCD')) + >>> df = pd.DataFrame(np.arange(12).reshape(3,4), + columns=['A', 'B', 'C', 'D']) >>> df - A B C D - 0 1 2 3 4 - 1 5 6 7 8 - 2 9 1 2 3 - 3 4 5 6 7 + A B C D + 0 0 1 2 3 + 1 4 5 6 7 + 2 8 9 10 11 + + Drop columns + + >>> df.drop(['B', 'C'], axis=1) + A D + 0 0 3 + 1 4 7 + 2 8 11 + + >>> df.drop(columns=['B', 'C']) + A D + 0 0 3 + 1 4 7 + 2 8 11 Drop a row by index >>> df.drop([0, 1]) - A B C D - 2 9 1 2 3 - 3 4 5 6 7 + A B C D + 2 8 9 10 11 - Drop columns + Notes + ----- + Specifying both `labels` and `index` or `columns` will raise a + ValueError. - >>> df.drop(['A', 'B'], axis=1) - C D - 0 3 4 - 1 7 8 - 2 2 3 - 3 6 7 """ inplace = validate_bool_kwarg(inplace, 'inplace') + + if labels is not None: + if index is not None or columns is not None: + raise ValueError("Cannot specify both 'labels' and " + "'index'/'columns'") + axis_name = self._get_axis_name(axis) + axes = {axis_name: labels} + elif index is not None or columns is not None: + axes, _ = self._construct_axes_from_arguments((index, columns), {}) + else: + raise ValueError("Need to specify at least one of 'labels', " + "'index' or 'columns'") + + obj = self + + for axis, labels in axes.items(): + if labels is not None: + obj = obj._drop_axis(labels, axis, level=level, errors=errors) + + if inplace: + self._update_inplace(obj) + else: + return obj + + def _drop_axis(self, labels, axis, level=None, errors='raise'): + """ + Drop labels from specified axis. Used in the ``drop`` method + internally. + + Parameters + ---------- + labels : single label or list-like + axis : int or axis name + level : int or level name, default None + For MultiIndex + errors : {'ignore', 'raise'}, default 'raise' + If 'ignore', suppress error and existing labels are dropped. + + """ axis = self._get_axis_number(axis) axis_name = self._get_axis_name(axis) axis, axis_ = self._get_axis(axis), axis @@ -2416,10 +2469,7 @@ def drop(self, labels, axis=0, level=None, inplace=False, errors='raise'): result = self.loc[tuple(slicer)] - if inplace: - self._update_inplace(result) - else: - return result + return result def _update_inplace(self, result, verify_is_copy=True): """ diff --git a/pandas/tests/frame/test_axis_select_reindex.py b/pandas/tests/frame/test_axis_select_reindex.py index e76869bf6712b..fb9b8c2ed7aff 100644 --- a/pandas/tests/frame/test_axis_select_reindex.py +++ b/pandas/tests/frame/test_axis_select_reindex.py @@ -146,6 +146,41 @@ def test_drop_multiindex_not_lexsorted(self): tm.assert_frame_equal(result, expected) + def test_drop_api_equivalence(self): + # equivalence of the labels/axis and index/columns API's (GH12392) + df = DataFrame([[1, 2, 3], [3, 4, 5], [5, 6, 7]], + index=['a', 'b', 'c'], + columns=['d', 'e', 'f']) + + res1 = df.drop('a') + res2 = df.drop(index='a') + tm.assert_frame_equal(res1, res2) + + res1 = df.drop('d', 1) + res2 = df.drop(columns='d') + tm.assert_frame_equal(res1, res2) + + res1 = df.drop(labels='e', axis=1) + res2 = df.drop(columns='e') + tm.assert_frame_equal(res1, res2) + + res1 = df.drop(['a'], axis=0) + res2 = df.drop(index=['a']) + tm.assert_frame_equal(res1, res2) + + res1 = df.drop(['a'], axis=0).drop(['d'], axis=1) + res2 = df.drop(index=['a'], columns=['d']) + tm.assert_frame_equal(res1, res2) + + with pytest.raises(ValueError): + df.drop(labels='a', index='b') + + with pytest.raises(ValueError): + df.drop(labels='a', columns='b') + + with pytest.raises(ValueError): + df.drop(axis=1) + def test_merge_join_different_levels(self): # GH 9455