covid-19-discord-bot/lib/covid_data_lib.py

#!/usr/bin/python3

from lib.covidData import covidData


def set_data_dict(selection, import_data):
    columns = [i for i in covidData.__dict__.keys() if i[:1] != '_']
    formatted_import_data = [d.strip().replace(
        ',', '').replace('+', '') for d in import_data]
    output_dict = {
        'selection': selection.upper(),
        'selection_original': selection
    }
    for i in range(2, len(columns)):
        if(formatted_import_data[i - 1] and not formatted_import_data[i - 1] == ''):
            if(i == len(columns) - 1):
                output_dict[columns[i]] = formatted_import_data[i - 1]
            else:
                output_dict[columns[i]] = int(formatted_import_data[i - 1])
    return output_dict


def get_data_dict(query_data):
    columns = [i for i in covidData.__dict__.keys() if i[:1] != '_']
    output_dict = {
        'selection': query_data.selection,
        'selection_original': query_data.selection_original
    }
    for i in range(2, len(columns)):
        if(i == len(columns) - 1):
            output_dict[columns[i]] = query_data.__getattribute__(
                columns[i])
        else:
            output_dict[columns[i]] = int(
                query_data.__getattribute__(columns[i]))
    return output_dict


def set_data(session, selection, import_data):
    new_data = covidData()
    formatted_data_dict = set_data_dict(selection, import_data)
    try:
        for key in formatted_data_dict:
            new_data.__setattr__(key, formatted_data_dict[key])
    except Exception as e:
        print(e)
    session.merge(new_data)
    session.commit()


def get_data(session, selection):
    print('Formatting data.')
    all_data_query = session.query(covidData).filter(
        covidData.selection == selection).all()
    return get_data_dict(all_data_query[0])


def get_top_n_rows(session, num):
    print(f'Getting top {num} rows.')
    top_n_rows = session.query(covidData).order_by(
        covidData.total_cases.desc()).limit(num).all()
    return [get_data_dict(n) for n in top_n_rows]