preprocess/utils/common/data.py

import pickle
import pandas as pd
import os.path

DATA_DIR = "_data"

def get_data_dir():
    return DATA_DIR


def set_data_dir(data_dir):
    global DATA_DIR
    DATA_DIR = data_dir


def get_products_dir():
    products_dir = os.path.join(get_data_dir(), "products")
    if not os.path.exists(products_dir):
        os.makedirs(products_dir)

    return products_dir


def get_results_dir():
    results_dir = os.path.join(get_data_dir(), "results")
    if not os.path.exists(results_dir):
        os.makedirs(results_dir)

    return results_dir

def get_temp_dir():
    temp_dir = os.path.join(get_data_dir(), "temp")
    if not os.path.exists(temp_dir):
        os.makedirs(temp_dir)

    return temp_dir


def load_df_from_file(filepath):
    if not os.path.isfile(filepath):
        return None
    return pd.read_pickle(filepath)


def save_df_to_file(df, filepath, save_csv_copy = False):
    df.to_pickle(filepath)

    if save_csv_copy:
        fullfn_csv = filepath + '.csv'
        df.to_csv(fullfn_csv, encoding="utf-8")


def load_dict_from_file(filepath):
    if not os.path.isfile(filepath):
        return None

    with open(filepath, 'rb') as fn:
        str = fn.read()
        return pickle.loads(str)


def save_dict_to_file(dict, filepath):
    with open(filepath, 'wb') as fn:
        str = pickle.dumps(dict)
        fn.write(str)