perfect-postcode/pipeline/journey_times/results.py

36 lines
962 B
Python

from pathlib import Path
import polars as pl
from .config import OUTPUT_DIR
from .models import JourneyResult
def results_to_dataframe(results: list[JourneyResult]) -> pl.DataFrame:
return pl.DataFrame(
[
{
"postcode": r.postcode,
"public_transport_easy_minutes": r.public_transport_easy_minutes,
"public_transport_quick_minutes": r.public_transport_quick_minutes,
"cycling_minutes": r.cycling_minutes,
"error": r.error,
}
for r in results
]
)
def save_results(
results: pl.DataFrame,
destination_name: str,
output_dir: Path | None = None,
) -> Path:
if output_dir is None:
output_dir = OUTPUT_DIR
safe_name = destination_name.lower().replace(" ", "-")
parquet_path = output_dir / f"journey_times_{safe_name}.parquet"
results.write_parquet(parquet_path)
return parquet_path