latecomers/latecomers/serializer.py

# -*- encoding: utf-8 -*-
import logging
import os.path
from tempfile import NamedTemporaryFile
import typing as T

from latecomers.parse import Details
from latecomers.helpers import logit, get_date

import pandas as pd


logger = logging.getLogger(__name__)


@logit(logger)
def to_excel(data: T.List[Details], dst: T.Optional[T.Text] = None) -> bytes:
    """
    This function takes the list of parsed rows as input and returns
    the bytes corresponding to the excel file derived from such lines.
    """

    mapping = {
        "th_arrival": "Arrivo teorico",
        "real_arrival": "Arrivo reale",
        "code": "Codice volo",
        "origin": "Aeroporto di partenza",
        "status": "Stato",
        "fr24_landing_time": "Ora atterraggio (FlightRadar24)",
    }
    df = pd.DataFrame(data, columns=mapping)
    df["status"] = df["status"].map(lambda x: x.value)
    df.set_index("th_arrival", inplace=True)
    df.rename(columns=mapping, inplace=True)
    df = df.rename_axis(index=mapping["th_arrival"])

    with NamedTemporaryFile() as tmp:
        df.to_excel(tmp)
        tmp.seek(0)
        content = tmp.read()

    if dst:
        filepath = os.path.join(dst, f"{get_date()}.xlsx")
        try:
            with open(filepath, "wb") as out:
                out.write(content)
        except Exception as e:
            logger.warning(f"Cannot save to path '{filepath}': {e}")

    return content