diff --git a/chainladder/core/pandas.py b/chainladder/core/pandas.py index 9d51741c..7d70cd62 100644 --- a/chainladder/core/pandas.py +++ b/chainladder/core/pandas.py @@ -354,7 +354,7 @@ def append(self, other): def rename( self, axis: Literal['index', 'columns', 'origin', 'development'] | int, - value: list | str + value: list | str | dict ) -> Self: """Alter axes labels. @@ -364,29 +364,40 @@ def rename( A value of 0 <= axis <= 4 corresponding to axes 'index', 'columns', 'origin', 'development' respectively. Both the int and str representation can be used. - value: list or str + value: list or str or dict List of new labels to be assigned to the axis. List must be of - same length of the specified axis. + same length of the specified axis. Can also be a dictionary for renaming columns Returns ------- Triangle with relabeled axis. """ - value = [value] if type(value) is str else value - if axis == "index" or axis == 0: - self.index = value - elif axis == "columns" or axis == 1: - self.columns = value - elif axis == "origin" or axis == 2: - self.origin = value - elif axis == "development" or axis == 3: - self.development = value + + if type(value) is dict: + if axis == "columns" or axis == 1: + full_dict = dict(zip(self.columns.values,self.columns.values)) + full_dict.update(value) + self.columns = self.columns.map(full_dict) + else: + raise ValueError( + "Invalid value provided to the 'value' parameter. Accepted values for index, origin, and development axes are a str or a list" + ) 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." - ) + value = [value] if type(value) is str else value + if axis == "index" or axis == 0: + self.index = value + elif axis == "columns" or axis == 1: + self.columns = value + elif axis == "origin" or axis == 2: + self.origin = value + 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 8070ee6e..e4803f5a 100644 --- a/chainladder/core/tests/test_triangle.py +++ b/chainladder/core/tests/test_triangle.py @@ -101,7 +101,20 @@ def test_rename_columns(genins, clrd) -> None: # Test the cascading of rename to triangle.columns_label. assert genins.columns_label == ['foo'] + genins.rename('columns',{'foo':'newfoo'}) + + assert genins.columns.to_list() == ['newfoo'] + + genins.rename('columns',{'foo':'newnewfoo'}) + + assert genins.columns.to_list() == ['newfoo'] + + def test_rename_exception(genins, clrd) -> None: + # Test incorrect value argument - misspelling of string. + with pytest.raises(ValueError): + genins.rename('origin', {'oldfoo':'foo'}) + # Test incorrect axis argument - misspelling of string. with pytest.raises(ValueError): genins.rename('colunms', 'foo')