Format python

pipeline/transform/join_epc_pp.py

@@ -9,79 +9,108 @@ pl.Config.set_tbl_cols(-1)
 
 def main():
     parser = argparse.ArgumentParser(description="Fuzzy join EPC and Price Paid data")
-    parser.add_argument("--epc", type=Path, required=True, help="EPC certificates CSV file")
-    parser.add_argument("--price-paid", type=Path, required=True, help="Price paid parquet file")
-    parser.add_argument("--output", type=Path, required=True, help="Output parquet file path")
+    parser.add_argument(
+        "--epc", type=Path, required=True, help="EPC certificates CSV file"
+    )
+    parser.add_argument(
+        "--price-paid", type=Path, required=True, help="Price paid parquet file"
+    )
+    parser.add_argument(
+        "--output", type=Path, required=True, help="Output parquet file path"
+    )
     args = parser.parse_args()
 
-    epc = pl.scan_csv(args.epc).select(
-        pl.col('ADDRESS').alias('epc_address'),
-        'POSTCODE',
-        'CURRENT_ENERGY_RATING',
-        'POTENTIAL_ENERGY_RATING',
-        pl.col('PROPERTY_TYPE').alias('epc_property_type'),
-        'BUILT_FORM',
-        'INSPECTION_DATE',
-        'TOTAL_FLOOR_AREA',
-        'NUMBER_HABITABLE_ROOMS',
-        'FLOOR_HEIGHT',
-        'CONSTRUCTION_AGE_BAND'
-    ).filter(pl.col('epc_address').is_not_null()).sort('INSPECTION_DATE', descending=True).group_by('epc_address', 'POSTCODE').first()
-
+    epc = (
+        pl.scan_csv(args.epc)
+        .select(
+            pl.col("ADDRESS").alias("epc_address"),
+            "POSTCODE",
+            "CURRENT_ENERGY_RATING",
+            "POTENTIAL_ENERGY_RATING",
+            pl.col("PROPERTY_TYPE").alias("epc_property_type"),
+            "BUILT_FORM",
+            "INSPECTION_DATE",
+            "TOTAL_FLOOR_AREA",
+            "NUMBER_HABITABLE_ROOMS",
+            "FLOOR_HEIGHT",
+            "CONSTRUCTION_AGE_BAND",
+        )
+        .filter(pl.col("epc_address").is_not_null())
+        .sort("INSPECTION_DATE", descending=True)
+        .group_by("epc_address", "POSTCODE")
+        .first()
+    )
 
     print("EPC dataset")
     print(epc.head().collect())
 
     # https://www.gov.uk/guidance/about-the-price-paid-data
-    property_type_map = {"D": "Detached", "S": "Semi-Detached", "T": "Terraced", "F": "Flats/Maisonettes", "O": "Other"}
+    property_type_map = {
+        "D": "Detached",
+        "S": "Semi-Detached",
+        "T": "Terraced",
+        "F": "Flats/Maisonettes",
+        "O": "Other",
+    }
     duration_map = {"F": "Freehold", "L": "Leasehold"}
 
-    price_paid = (pl.scan_parquet(args.price_paid).select(
-        "price",
-        "date_of_transfer",
-        pl.col('property_type').alias("pp_property_type").replace(property_type_map),
-        "postcode",
-        'paon',
-        'saon',
-        'street',
-        'locality',
-        'town_city',
-        pl.col('duration').replace(duration_map)
-    )
-    .filter(pl.col('pp_property_type') != 'Other').with_columns(
-            pl.concat_str(
-                [pl.col('saon'), pl.col('paon'), pl.col('street')],
-                separator=' ',
-                ignore_nulls=True,
-            ).alias('pp_address'),
+    price_paid = (
+        pl.scan_parquet(args.price_paid)
+        .select(
+            "price",
+            "date_of_transfer",
+            pl.col("property_type")
+            .alias("pp_property_type")
+            .replace(property_type_map),
+            "postcode",
+            "paon",
+            "saon",
+            "street",
+            "locality",
+            "town_city",
+            pl.col("duration").replace(duration_map),
         )
-        .sort('date_of_transfer')
-        .group_by('pp_address', 'postcode', maintain_order=True)
+        .filter(pl.col("pp_property_type") != "Other")
+        .with_columns(
+            pl.concat_str(
+                [pl.col("saon"), pl.col("paon"), pl.col("street")],
+                separator=" ",
+                ignore_nulls=True,
+            ).alias("pp_address"),
+        )
+        .sort("date_of_transfer")
+        .group_by("pp_address", "postcode", maintain_order=True)
         .agg(
             pl.struct(
-                pl.col('date_of_transfer').dt.year().alias('year'),
-                'price',
-            ).alias('historical_prices'),
-            pl.col('pp_property_type').last(),
-            pl.col('duration').last(),
-            pl.col('price').last().alias('latest_price'),
-            pl.col('date_of_transfer').last(),
+                pl.col("date_of_transfer").dt.year().alias("year"),
+                "price",
+            ).alias("historical_prices"),
+            pl.col("pp_property_type").last(),
+            pl.col("duration").last(),
+            pl.col("price").last().alias("latest_price"),
+            pl.col("date_of_transfer").last(),
         )
-    ).filter(pl.col('pp_address').is_not_null())
+    ).filter(pl.col("pp_address").is_not_null())
 
     print("Price paid dataset")
     print(price_paid.head().collect())
 
-    joined = fuzzy_join_on_postcode(
-        left=price_paid,
-        right=epc,
-        left_address_col='pp_address',
-        right_address_col='epc_address',
-        left_postcode_col='postcode',
-        right_postcode_col='POSTCODE',
-    ).drop('POSTCODE').collect()
+    joined = (
+        fuzzy_join_on_postcode(
+            left=price_paid,
+            right=epc,
+            left_address_col="pp_address",
+            right_address_col="epc_address",
+            left_postcode_col="postcode",
+            right_postcode_col="POSTCODE",
+        )
+        .drop("POSTCODE")
+        .collect()
+    )
 
-    matched = joined.filter(pl.col('epc_address').is_not_null() & pl.col('pp_address').is_not_null())
+    matched = joined.filter(
+        pl.col("epc_address").is_not_null() & pl.col("pp_address").is_not_null()
+    )
     total = joined.height
     print(f"Unique properties: {total}")
     print(f"Matched: {matched.height} ({100 * matched.height / total:.1f}%)")