format

PSLmodels · Dec 22, 2024 · 4b9e4d7 · 4b9e4d7
1 parent 55c867c
commit 4b9e4d7
Show file tree

Hide file tree

Showing 2 changed files with 48 additions and 29 deletions.
diff --git 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
@@ -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)