utils.py

"""
Helper functions for our project

... Requires logging already set up!
"""

# general
# import logging
import pdb

# data science
import pandas as pd
import json

# settings
from settings import *


def read_csv_to_df(filename, converters=None):
    """
    Read a given csv file to data frame. Exit program if anything goes
    wrong.

    input: filename, columns
    output: dataframe
    """
    try:
        df = pd.read_csv(filename, converters=converters)
        return df
    except Exception as exc:
        # logging.error(f'Error reading csv file {filename}')
        print(exc)
        exit()


def read_pkl_to_df(filename):
    """
    Read pickle file to pandas dataframe
    """
    try:
        df = pd.read_pickle(filename)
        return df
    except Exception as exc:
        # logging.error(f'Error reading pickle {filename}')
        print(exc)


def save_data_to_file(df, filename):
    """
    Save data frame to pickle or csv
    """
    try:
        if filename.endswith('pkl'):
            df.to_pickle(filename)
            # logging.info(f'Data saved to pickle {filename}')
        elif filename.endswith('csv'):
            df.to_csv(filename, index=False)
            # logging.info(f'Data saved to csv {filename}')

    except Exception as exc:
        # logging.warning(f'Could not write file {filename}')
        print(exc)


if __name__ == '__main__':
    """
    Tryouts
    """
    pass