wrongmove/csv_exporter.py

from pathlib import Path
import pandas as pd
from models.listing import QueryParameters
from repositories.listing_repository import ListingRepository


async def export_to_csv(
    repository: ListingRepository,
    output_file: Path,
    query_parameters: QueryParameters | None = None,
) -> None:
    listings = await repository.get_listings(query_parameters=query_parameters)
    ds = [listing.__dict__ for listing in listings]
    df = pd.DataFrame(ds)

    # read decisions on file
    decisions_path = Path("data/decisions.json")
    if decisions_path.exists():
        decisions = pd.read_json(decisions_path)
        df.loc[:, "decision"] = df.id.apply(lambda x: decisions.get(x))

    # remove _sa_instance_state column
    drop_columns = ["_sa_instance_state", "additional_info"]
    df = df.drop(columns=drop_columns)

    # Ensure columns exist with NaN defaults for clean CSV output
    for col in ("service_charge", "lease_left", "square_meters"):
        if col not in df.columns:
            df.loc[:, col] = float("nan")

    # Replace -1 sentinel values with NaN
    df.loc[:, "square_meters"] = df.square_meters.replace({-1: float("nan")})

    # Add price per sqm column (guard against zero/missing square_meters)
    df.loc[:, "price_per_sqm"] = df.apply(
        lambda row: round(row.price / row.square_meters, 2)
        if row.square_meters and row.square_meters > 0
        else None,
        axis=1,
    )

    df = df.sort_values(by=["price_per_sqm"], ascending=True)
    df.to_csv(str(output_file), index=False)