Update grid_loss_responsible.py

Energinet-DataHub · Sep 14, 2023 · d2507e1 · d2507e1
1 parent e98a1df
commit d2507e1
Showing 1 changed file with 47 additions and 19 deletions.
diff --git a/source/databricks/calculation-engine/package/calculation_input/grid_loss_responsible.py b/source/databricks/calculation-engine/package/calculation_input/grid_loss_responsible.py
@@ -26,6 +26,7 @@
 
 DEFAULT_FROM_TIME = datetime.strptime("2000-01-01T23:00:00+0000", "%Y-%m-%dT%H:%M:%S%z")
 
+# fmt: off
 GRID_AREA_RESPONSIBLE = [
     # u-001 and t-001
     ('571313180480500149', 804, DEFAULT_FROM_TIME, None, MeteringPointType.PRODUCTION.value, '8100000000108'),
@@ -44,17 +45,25 @@
     ('571313174115776740', 791, DEFAULT_FROM_TIME, None, MeteringPointType.CONSUMPTION.value, '5790000701414'),
     ('571313179100348995', 791, DEFAULT_FROM_TIME, None, MeteringPointType.PRODUCTION.value, '5790002437717'),
 ]
+# fmt: on
 
 
 def get_grid_loss_responsible(grid_areas: list[str]) -> DataFrame:
     grid_loss_responsible_df = _get_all_grid_loss_responsible()
 
     grid_loss_responsible_df = grid_loss_responsible_df.withColumn(
-        Colname.is_positive_grid_loss_responsible, when(col(Colname.metering_point_type) == MeteringPointType.CONSUMPTION.value, True)
-        .otherwise(False))
+        Colname.is_positive_grid_loss_responsible,
+        when(
+            col(Colname.metering_point_type) == MeteringPointType.CONSUMPTION.value,
+            True,
+        ).otherwise(False),
+    )
     grid_loss_responsible_df = grid_loss_responsible_df.withColumn(
-        Colname.is_negative_grid_loss_responsible, when(col(Colname.metering_point_type) == MeteringPointType.PRODUCTION.value, True)
-        .otherwise(False))
+        Colname.is_negative_grid_loss_responsible,
+        when(
+            col(Colname.metering_point_type) == MeteringPointType.PRODUCTION.value, True
+        ).otherwise(False),
+    )
 
     grid_loss_responsible_df = grid_loss_responsible_df.select(
         col(Colname.metering_point_id),
@@ -72,25 +81,44 @@ def get_grid_loss_responsible(grid_areas: list[str]) -> DataFrame:
     return grid_loss_responsible_df
 
 
-def _throw_if_no_grid_loss_responsible(grid_areas: list[str], grid_loss_responsible_df: DataFrame) -> None:
-
+def _throw_if_no_grid_loss_responsible(
+    grid_areas: list[str], grid_loss_responsible_df: DataFrame
+) -> None:
     for grid_area in grid_areas:
-        current_grid_area_responsible = grid_loss_responsible_df.filter(col(Colname.grid_area) == grid_area)
-        if (current_grid_area_responsible.filter(col(Colname.is_negative_grid_loss_responsible)).count() == 0):
-            raise ValueError(f"No responsible for negative grid loss found for grid area {grid_area}")
-        if (current_grid_area_responsible.filter(col(Colname.is_positive_grid_loss_responsible)).count() == 0):
-            raise ValueError(f"No responsible for positive grid loss found for grid area {grid_area}")
+        current_grid_area_responsible = grid_loss_responsible_df.filter(
+            col(Colname.grid_area) == grid_area
+        )
+        if (
+            current_grid_area_responsible.filter(
+                col(Colname.is_negative_grid_loss_responsible)
+            ).count()
+            == 0
+        ):
+            raise ValueError(
+                f"No responsible for negative grid loss found for grid area {grid_area}"
+            )
+        if (
+            current_grid_area_responsible.filter(
+                col(Colname.is_positive_grid_loss_responsible)
+            ).count()
+            == 0
+        ):
+            raise ValueError(
+                f"No responsible for positive grid loss found for grid area {grid_area}"
+            )
 
 
 def _get_all_grid_loss_responsible() -> DataFrame:
-    schema = StructType([
-        StructField(Colname.metering_point_id, StringType(), nullable=False),
-        StructField(Colname.grid_area, StringType(), nullable=False),
-        StructField(Colname.from_date, TimestampType(), nullable=False),
-        StructField(Colname.to_date, TimestampType(), nullable=True),
-        StructField(Colname.metering_point_type, StringType(), nullable=False),
-        StructField(Colname.energy_supplier_id, StringType(), nullable=False)
-    ])
+    schema = StructType(
+        [
+            StructField(Colname.metering_point_id, StringType(), nullable=False),
+            StructField(Colname.grid_area, StringType(), nullable=False),
+            StructField(Colname.from_date, TimestampType(), nullable=False),
+            StructField(Colname.to_date, TimestampType(), nullable=True),
+            StructField(Colname.metering_point_type, StringType(), nullable=False),
+            StructField(Colname.energy_supplier_id, StringType(), nullable=False),
+        ]
+    )
 
     spark = SparkSession.builder.getOrCreate()