diff --git a/chainladder/core/pandas.py b/chainladder/core/pandas.py index 26a9e6ee..8c90e090 100644 --- a/chainladder/core/pandas.py +++ b/chainladder/core/pandas.py @@ -20,7 +20,11 @@ from collections.abc import Callable from numpy.typing import ArrayLike from types import ModuleType - from typing import Type + from typing import ( + Literal, + Self, + Type + ) @@ -308,7 +312,11 @@ def append(self, other): return concat((self, other), 0) - def rename(self, axis, value): + def rename( + self, + axis: Literal['index', 'columns', 'origin', 'development'] | int, + value: list | str + ) -> Self: """Alter axes labels. Parameters @@ -328,12 +336,18 @@ def rename(self, axis, value): value = [value] if type(value) is str else value if axis == "index" or axis == 0: self.index = value - if axis == "columns" or axis == 1: + elif axis == "columns" or axis == 1: self.columns = value - if axis == "origin" or axis == 2: + elif axis == "origin" or axis == 2: self.origin = value - if axis == "development" or axis == 3: + elif axis == "development" or axis == 3: self.development = value + else: + raise ValueError( + "Invalid value provided to the 'axis' parameter. Accepted values are a string of 'index', " + "'columns', 'origin', or 'development', or an integer in the interval [0, 4] specifying the" + " axis to be modified." + ) return self def astype(self, dtype, inplace=True): diff --git a/chainladder/core/tests/test_triangle.py b/chainladder/core/tests/test_triangle.py index 47cc59df..8ff5d638 100644 --- a/chainladder/core/tests/test_triangle.py +++ b/chainladder/core/tests/test_triangle.py @@ -89,6 +89,30 @@ def test_append(raa): raa.append(raa2).sum() == raa * 2 assert raa.append(raa2).sum() == 2 * raa +def test_rename_columns(genins, clrd) -> None: + """ + Test the renaming of triangle columns. + """ + # Scalar case - single column triangle. + genins.rename('columns', 'foo') + + assert genins.columns.to_list() == ['foo'] + + # Test the cascading of rename to triangle.columns_label. + assert genins.columns_label == ['foo'] + +def test_rename_exception(genins, clrd) -> None: + # Test incorrect axis argument - misspelling of string. + with pytest.raises(ValueError): + genins.rename('colunms', 'foo') + + # Test incorrect axis integer. + with pytest.raises(ValueError): + genins.rename(4, 'foo') + + # Test incorrect number of columns. + with pytest.raises(ValueError): + clrd.rename('columns', ['foo']) def test_assign_existing_col(qtr): out = qtr.copy() diff --git a/chainladder/core/triangle.py b/chainladder/core/triangle.py index 041f9111..09a3361b 100644 --- a/chainladder/core/triangle.py +++ b/chainladder/core/triangle.py @@ -149,7 +149,6 @@ def __init__( ) # Store dimension metadata. - self.columns_label: list = columns self.origin_label: list = origin # Handle any ultimate vectors in triangles separately. @@ -401,6 +400,10 @@ def columns(self, value): self.vdims = np.array(self.vdims) self._set_slicers() + @property + def columns_label(self) -> list: + return self.columns.to_list() + @property def origin(self): if self.is_pattern and len(self.odims) == 1: