diff --git a/ccc/parameters.py b/ccc/parameters.py index 52ec562a..abc3f887 100644 --- a/ccc/parameters.py +++ b/ccc/parameters.py @@ -343,11 +343,11 @@ def to_df(self): # Prepare a dictionary to hold the row data row_data = {} # Extract 'value' dictionary contents - row_data.update(item['value']) + row_data.update(item["value"]) # Add the year - row_data['year'] = item['year'] + row_data["year"] = item["year"] # Add the asset code - row_data['BEA_code'] = key + row_data["BEA_code"] = key # Append the row data and index key data.append(row_data) @@ -358,21 +358,26 @@ def to_df(self): def expanded_df(self): df = self.to_df() - unique_bea_codes = df['BEA_code'].unique() + unique_bea_codes = df["BEA_code"].unique() years = self.label_grid["year"] # Create all combinations of years and BEA codes combinations = list(itertools.product(years, unique_bea_codes)) # Create a new DataFrame with these combinations - expanded_df = pd.DataFrame(combinations, columns=['year', 'BEA_code']) + expanded_df = pd.DataFrame(combinations, columns=["year", "BEA_code"]) # Merge with the original DataFrame to preserve known values - expanded_df = expanded_df.merge(df, on=['year', 'BEA_code'], how='left') + expanded_df = expanded_df.merge( + df, on=["year", "BEA_code"], how="left" + ) # Sort the DataFrame to ensure we can forward fill from the last known year - expanded_df = expanded_df.sort_values(['BEA_code', 'year']) + expanded_df = expanded_df.sort_values(["BEA_code", "year"]) # Group by BEA_code and forward fill missing values - expanded_df = expanded_df[['life', 'method', 'system', 'year', 'BEA_code']].groupby('BEA_code').apply( - lambda group: group.ffill(), include_groups=False - ).reset_index() - column_order = ['life', 'method', 'system', 'year', 'BEA_code'] + expanded_df = ( + expanded_df[["life", "method", "system", "year", "BEA_code"]] + .groupby("BEA_code") + .apply(lambda group: group.ffill(), include_groups=False) + .reset_index() + ) + column_order = ["life", "method", "system", "year", "BEA_code"] expanded_df = expanded_df[column_order] return expanded_df diff --git a/ccc/tests/test_parameters.py b/ccc/tests/test_parameters.py index dd27b9ab..8841f5a5 100644 --- a/ccc/tests/test_parameters.py +++ b/ccc/tests/test_parameters.py @@ -174,15 +174,17 @@ def test_update_depreciation_params_with_dict(): } dp = DepreciationParams() new_dp_dict = { - "ENS2": [{ - "year": 2020, - "value": {"life": 5.0, "method": "Expensing", "system": "GDS"} - }] + "ENS2": [ + { + "year": 2020, + "value": {"life": 5.0, "method": "Expensing", "system": "GDS"}, } + ] + } dp.adjust(new_dp_dict) - test_result = dp.select_eq( - param="ENS2", strict=False, year=2020 - )[0]["value"] + test_result = dp.select_eq(param="ENS2", strict=False, year=2020)[0][ + "value" + ] print("test", test_result) print("expected", expected_result) assert test_result == expected_result @@ -201,9 +203,9 @@ def test_update_depreciation_params_with_json(): "value": {"life": 5, "method": "Expensing", "system": "GDS"}}]} """ dp.adjust(new_dp_json) - test_result = dp.select_eq( - param="ENS2", strict=False, year=2020 - )[0]["value"] + test_result = dp.select_eq(param="ENS2", strict=False, year=2020)[0][ + "value" + ] assert test_result == expected_result @@ -235,11 +237,17 @@ def test_update_depreciation_bad_revision(): """ dp = DepreciationParams() new_dp_dict = { - "ENS2": [{ - "year": 2020, - "value": {"life": 5.0, "method": "Expensing2", "system": "GDS"} - }] + "ENS2": [ + { + "year": 2020, + "value": { + "life": 5.0, + "method": "Expensing2", + "system": "GDS", + }, } + ] + } with pytest.raises(Exception): assert dp.adjust(new_dp_dict) @@ -250,10 +258,16 @@ def test_update_depreciation_bad_revision2(): """ dp = DepreciationParams() new_dp_dict = { - "ENS2": [{ - "year": 2020, - "value": {"life": 105.0, "method": "Expensing2", "system": "GDS"} - }] + "ENS2": [ + { + "year": 2020, + "value": { + "life": 105.0, + "method": "Expensing2", + "system": "GDS", + }, } + ] + } with pytest.raises(Exception): assert dp.adjust(new_dp_dict)