SPlit up

pipeline/download/ethnicity.py

@@ -16,12 +16,12 @@ License: Open Government Licence v3.0
 """
 
 import argparse
-from io import BytesIO
 from pathlib import Path
 
-import httpx
 import polars as pl
 
+from pipeline.utils import ENGLAND_LSOA_COUNT_2021, download_nomis_csv
+
 pl.Config.set_tbl_cols(-1)
 
 # NOMIS API: Census 2021 TS021 (ethnic group, 20 categories) by LSOA 2021
@@ -35,7 +35,6 @@ BASE_URL = (
     "&measures=20100"
     "&select=GEOGRAPHY_CODE,C2021_ETH_20_NAME,OBS_VALUE"
 )
-PAGE_SIZE = 25000
 
 # Map the 19 detailed NOMIS C2021_ETH_20 leaf categories to our 7 output groups.
 # The Asian split:
@@ -150,24 +149,7 @@ def _ethnicity_percentages(df: pl.DataFrame) -> pl.DataFrame:
 
 def download_and_convert(output_path: Path) -> None:
     print("Downloading Census 2021 ethnic group (TS021) by LSOA from NOMIS...")
-    frames = []
-    offset = 0
-    while True:
-        url = f"{BASE_URL}&recordoffset={offset}"
-        response = httpx.get(url, follow_redirects=True, timeout=120)
-        response.raise_for_status()
-        if len(response.content) == 0:
-            break
-        chunk = pl.read_csv(BytesIO(response.content))
-        if chunk.height == 0:
-            break
-        frames.append(chunk)
-        print(f"  Fetched {chunk.height} rows (offset={offset})")
-        if chunk.height < PAGE_SIZE:
-            break
-        offset += PAGE_SIZE
-
-    df = pl.concat(frames)
+    df = download_nomis_csv(BASE_URL)
     print(f"Total rows: {df.height}")
 
     # Filter to England only (E-prefixed LSOA codes); the merge joins on the
@@ -177,6 +159,11 @@ def download_and_convert(output_path: Path) -> None:
     wide = _ethnicity_percentages(df)
 
     print(f"England LSOAs: {wide.height}")
+    if wide.height != ENGLAND_LSOA_COUNT_2021:
+        raise ValueError(
+            f"Expected {ENGLAND_LSOA_COUNT_2021} England LSOAs, "
+            f"got {wide.height}: truncated NOMIS download?"
+        )
     print(f"Columns: {wide.columns}")
 
     output_path.parent.mkdir(parents=True, exist_ok=True)