piket/piket_client/export.py

"""
Provides functionality for creating Export files.
"""
import csv
import os
import sys
from piket_client.model import Settlement, Export, ConsumptionType, Person


if len(sys.argv) > 1:
    eid = int(sys.argv[1])
    export = Export.get(eid)
else:
    export = Export.create()

filename_base = export.created_at.strftime("Export %Y-%m-%d")

# Determine which consumption types exist
consumption_types = [
    ConsumptionType.get(id)
    for id in {
        type_id
        for s in export.settlements
        for type_id in s["consumption_summary"].keys()
    }
]

consumption_types.sort(key=lambda x: x.name)

extended_settlements = [
    Settlement.get(id) for id in map(lambda x: x["settlement_id"], export.settlements)
]
extended_settlements.sort(key=lambda x: x.name)

person_ids = set()
for settlement in extended_settlements:
    for ctype in settlement.count_info.values():
        for person_id in ctype["counts"].keys():
            person_ids.add(person_id)

persons = [Person.get(person_id) for person_id in person_ids]
persons.sort(key=lambda p: p.name)

fieldnames = [
    "Naam",
    *map(lambda x: f"{x.name} ({x.settlement_id})", extended_settlements),
    "Totaal",
]

# Create overview table
with open(f"{filename_base} Overzicht.csv", "w") as outf:

    writer = csv.DictWriter(outf, fieldnames)

    writer.writeheader()

    for consumption_type in consumption_types:
        row = {"Naam": consumption_type.name}
        total = 0
        for settlement in extended_settlements:
            strid = str(consumption_type.consumption_type_id)
            info = settlement.consumption_summary.get(strid, {})

            cell = info.get("count", "")
            if cell:
                total += cell

            row[f"{settlement.name} ({settlement.settlement_id})"] = cell

        row["Totaal"] = total
        writer.writerow(row)

# Create table per consumption type
for c_type in consumption_types:
    strid = str(c_type.consumption_type_id)
    with open(f"{filename_base} {c_type.name}.csv", "w") as outf:
        writer = csv.DictWriter(outf, fieldnames)
        writer.writeheader()

        for person in persons:
            row = {"Naam": person.name}
            total = 0

            for settlement in extended_settlements:
                info = settlement.count_info.get(strid, {})
                countinfo = info.get("counts", {})
                pinfo = countinfo.get(str(person.person_id), {})

                cell = pinfo.get("count", "")
                if cell:
                    total += cell

                row[f"{settlement.name} ({settlement.settlement_id})"] = cell

            row["Totaal"] = total
            writer.writerow(row)