Add more datasets

pipeline/download/naptan.py

@@ -0,0 +1,57 @@
+"""Download NaPTAN data and extract railway/metro station POIs."""
+
+import argparse
+import io
+import urllib.request
+from pathlib import Path
+
+import polars as pl
+
+NAPTAN_CSV_URL = "https://naptan.api.dft.gov.uk/v1/access-nodes?dataFormat=csv"
+
+
+def download_naptan(output: Path) -> None:
+    output.parent.mkdir(parents=True, exist_ok=True)
+
+    print(f"Downloading NaPTAN data from {NAPTAN_CSV_URL}")
+    with urllib.request.urlopen(NAPTAN_CSV_URL) as resp:
+        raw = resp.read()
+
+    print(f"Downloaded {len(raw) / (1024 * 1024):.1f} MB")
+
+    df = (
+        pl.read_csv(io.BytesIO(raw), infer_schema_length=0)
+        .with_columns(
+            pl.col("Latitude").cast(pl.Float64, strict=False),
+            pl.col("Longitude").cast(pl.Float64, strict=False),
+        )
+        .drop_nulls(subset=["Latitude", "Longitude"])
+        .select(
+            pl.col("ATCOCode").alias("id"),
+            pl.col("CommonName").alias("name"),
+            pl.col("StopType").alias("category"),
+            pl.col("Latitude").alias("lat"),
+            pl.col("Longitude").alias("lng"),
+        )
+    )
+
+    df.write_parquet(output)
+    size_mb = output.stat().st_size / (1024 * 1024) 
+    print(f"Wrote {output} ({size_mb:.1f} MB, {len(df):,} stations)")
+
+    counts = df.group_by("category").len().sort("len", descending=True)
+    for row in counts.iter_rows(named=True):
+        print(f"  {row['category']}: {row['len']:,}")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Download NaPTAN station data")
+    parser.add_argument(
+        "--output", type=Path, required=True, help="Output parquet file path"
+    )
+    args = parser.parse_args()
+    download_naptan(args.output)
+
+
+if __name__ == "__main__":
+    main()