system-pri/backend/app/coordinator/utils.py

from typing import Generator, Any, List
from collections import defaultdict

import pandas as pd

from .exceptions import InvalidNameOrTypeHeaderException
from ..students.models import Student


def check_columns(df: pd.DataFrame) -> bool:
    headers = set(df.keys().values)
    columns = ['NAZWISKO', 'IMIE', 'INDEKS', 'PESEL', 'EMAIL']

    if len(headers - set(columns)) != 0:
        return False

    flag = True
    col_types = ['object', 'object', 'int', 'float64', 'object']

    for name, col_type in zip(columns, col_types):
        if not str(df.dtypes[name]).startswith(col_type):
            flag = False
            break

    return flag


def parse_csv(file, mode) -> Generator[Student, Any, None]:
    df = pd.read_csv(file)

    if not check_columns(df):
        raise InvalidNameOrTypeHeaderException

    students = (Student(last_name=dict(item.items())['NAZWISKO'],
                        first_name=dict(item.items())['IMIE'],
                        index=dict(item.items())['INDEKS'],
                        pesel=str(int(dict(item.items())['PESEL'])) if not pd.isna(
                            dict(item.items())['PESEL']) else None,
                        email=dict(item.items())['EMAIL'],
                        mode=mode)
                for _, item in df.iterrows())

    return students


def generate_csv(students: List[Student]) -> str:
    headers = ['PESEL', 'INDEKS', 'IMIE', 'NAZWISKO', 'EMAIL', 'CDYD_KOD', 'PRZ_KOD', 'TZAJ_KOD', 'GR_NR', 'PRG_KOD']
    data = [(student.pesel, student.index, student.first_name, student.last_name, student.email,
             student.group.cdyd_kod, student.group.prz_kod, student.group.tzaj_kod, student.group.project_supervisor_id,
             None) for student in students]
    dataframe = defaultdict(list)
    for row in data:
        for idx, item in enumerate(row):
            dataframe[headers[idx]].append(item)

    df = pd.DataFrame(dataframe)
    return df.to_csv(index=False)