diff --git a/chainladder/core/pandas.py b/chainladder/core/pandas.py index 35a83b1b1..d0f98eeb9 100644 --- a/chainladder/core/pandas.py +++ b/chainladder/core/pandas.py @@ -601,33 +601,100 @@ def drop( self, labels: str | int | list | None = None, axis: Literal["index", "columns", "origin", "development"] | int = 1, + index: str | int | list | None = None, + columns: str | int | list | None = None, + origin: str | int | list | None = None, + development: str | int | list | None = None, ) -> Triangle: """Drop specified labels from rows or columns. - Remove rows or columns by specifying label names and corresponding axis, - or by specifying directly index or column names. + Remove labels by specifying label names and corresponding axis, or by + specifying directly ``index``, ``columns``, ``origin``, or + ``development`` names. Parameters ----------- labels: str | int | list | None - Index or column labels to drop. + Index or column labels to drop. A single label or list-like. - axis: {0 or ‘index’, 1 or ‘columns’}, default 1 - Whether to drop labels from the index (0 or ‘index’) - or columns (1 or ‘columns’). + axis: {0 or ‘index’, 1 or ‘columns’, 2 or 'origin', 3 or 'development'}, default 1 + The axis to drop ``labels`` from. + + index: str | int | list | None + Alternative to ``axis=0``. Equivalent to ``labels, axis=0``. + + columns: str | int | list | None + Alternative to ``axis=1``. Equivalent to ``labels, axis=1``. + + origin: str | int | list | None + Alternative to ``axis=2``. Equivalent to ``labels, axis=2``. + + development: str | int | list | None + Alternative to ``axis=3``. Equivalent to ``labels, axis=3``. Returns ------- Triangle + Examples + -------- + + Drop a single column with the ``labels``/``axis`` form or the + ``columns`` alternative; the two are equivalent. + + .. testsetup:: + + import chainladder as cl + + .. testcode:: + + tri = cl.load_sample('clrd') + print(tri.columns.tolist()) + print(tri.drop(columns='CumPaidLoss').columns.tolist()) + + .. testoutput:: + + ['IncurLoss', 'CumPaidLoss', 'BulkLoss', 'EarnedPremDIR', 'EarnedPremCeded', 'EarnedPremNet'] + ['IncurLoss', 'BulkLoss', 'EarnedPremDIR', 'EarnedPremCeded', 'EarnedPremNet'] + + A list of labels can be dropped from an axis as well. + + .. testcode:: + + print(tri.drop(columns=['CumPaidLoss', 'IncurLoss']).columns.tolist()) + + .. testoutput:: + + ['BulkLoss', 'EarnedPremDIR', 'EarnedPremCeded', 'EarnedPremNet'] + """ - axis = self._get_axis(axis) - labels = [labels] if type(labels) is str else list(labels) - if axis == 1: - return self[[item for item in self.columns if item not in labels]] + alternatives = {0: index, 1: columns, 2: origin, 3: development} + if any(value is not None for value in alternatives.values()): + if labels is not None: + raise ValueError( + "Cannot specify both 'labels' and any of 'index', " + "'columns', 'origin', or 'development'." + ) + to_drop = { + ax: value for ax, value in alternatives.items() if value is not None + } else: - raise NotImplementedError("Triangle.drop() only implemented for column axis.") + to_drop = {self._get_axis(axis): labels} + result = self + for ax, ax_labels in to_drop.items(): + ax_labels = ( + [ax_labels] if np.isscalar(ax_labels) else list(ax_labels) + ) + if ax == 1: + result = result[ + [item for item in result.columns if item not in ax_labels] + ] + else: + raise NotImplementedError( + "Triangle.drop() only implemented for column axis." + ) + return result @property def T(self) -> DataFrame: # noqa: N802 diff --git a/chainladder/core/tests/test_triangle.py b/chainladder/core/tests/test_triangle.py index 1a0857e79..2a7ceb5f5 100644 --- a/chainladder/core/tests/test_triangle.py +++ b/chainladder/core/tests/test_triangle.py @@ -338,6 +338,52 @@ def test_drop_invalid_axis_raises(clrd): clrd.drop(labels="CumPaidLoss", axis="bogus") +def test_drop_columns_alternative(clrd): + """columns= should be equivalent to labels=..., axis=1.""" + result = clrd.drop(columns="CumPaidLoss") + assert "CumPaidLoss" not in result.columns + assert result == clrd.drop(labels="CumPaidLoss", axis=1) + + +def test_drop_columns_alternative_list(clrd): + """columns= should accept a list of labels.""" + result = clrd.drop(columns=["CumPaidLoss", "IncurLoss"]) + assert "CumPaidLoss" not in result.columns + assert "IncurLoss" not in result.columns + assert result == clrd.drop(labels=["CumPaidLoss", "IncurLoss"], axis=1) + + +def test_drop_labels_and_alternative_raises(clrd): + """Specifying labels together with an alternative should raise ValueError.""" + with pytest.raises(ValueError): + clrd.drop(labels="CumPaidLoss", columns="IncurLoss") + + +def test_drop_index_origin_development_alternatives_raise(clrd): + """index/origin/development alternatives route to unimplemented axes.""" + with pytest.raises(NotImplementedError): + clrd.drop(index="commauto") + with pytest.raises(NotImplementedError): + clrd.drop(origin="1995") + with pytest.raises(NotImplementedError): + clrd.drop(development="12") + + +def test_drop_integer_label_routes_to_axis(clrd): + """A bare int label should route to its axis, not raise TypeError.""" + with pytest.raises(NotImplementedError): + clrd.drop(development=12) + + +def test_drop_columns_alternative_index_like(clrd): + """columns= should accept list-likes such as pd.Index and ndarray.""" + labels = clrd.columns[:2] + result = clrd.drop(columns=labels) + assert all(label not in result.columns for label in labels) + assert result == clrd.drop(columns=list(labels)) + assert result == clrd.drop(columns=labels.values) + + def test_exposure_tri(): x = cl.load_sample("auto") x = x[x.development == 12] diff --git a/chainladder/core/typing.py b/chainladder/core/typing.py index 5f39b3332..ae0405193 100644 --- a/chainladder/core/typing.py +++ b/chainladder/core/typing.py @@ -97,7 +97,15 @@ def __len__(self) -> int: ... def get_array_module(self, arr: ArrayLike | None = None) -> ModuleType: ... def copy(self) -> Triangle: ... def set_backend(self, backend: str, inplace: bool = False, **kwargs) -> Triangle: ... - def drop(self, labels: str | int | list | None = None, axis: Literal["index", "columns", "origin", "development"] | int = 1) -> Triangle: ... + def drop( + self, + labels: str | int | list | None = None, + axis: Literal["index", "columns", "origin", "development"] | int = 1, + index: str | int | list | None = None, + columns: str | int | list | None = None, + origin: str | int | list | None = None, + development: str | int | list | None = None, + ) -> Triangle: ... def val_to_dev(self) -> Triangle: ... def _repr_format(self, origin_as_datetime: bool = False) -> pd.DataFrame: ... def _slice(self, key: pd.Series | np.ndarray, axis: Literal['ddims', 'odims']) -> Triangle: ...